modular / modular/modular

[Feature Request] Add new MAX ops `bitwise_and`, `bitwise_or` and `bitwise_xor`

Open
#5,300 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement gex-graph-api graph-api max
Dominant language
Mojo
Stars
29.8k
Forks
3.2k
PR merge metrics
No merged PRs in 30d

Description

What is your request?

torch functions:

Note that we could provide ops with those names but also if we want to follow torch's api, we'll want to implement it in TensorValue's __and__, __or__ and __xor__. The behavior in pytorch for those is to use the bitwise version:

In [3]: torch.tensor([1, 2], dtype=torch.uint8) |  torch.tensor([2, 4], dtype=torch.uint8)
Out[3]: tensor([3, 6], dtype=torch.uint8)
What is your motivation for this change?

I need it to implement the max backend for torch, and surprisingly it's used in gemma3 (or at least it's what torch aotautograd decomposes some ops into).

Any other details?

I could code them in torch-max-backend directly in Mojo, but then I end up with this bug: https://github.com/modular/modular/issues/5299

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the linked PyTorch bitwise_and, bitwise_or, and bitwise_xor documentation and inspect the MAX TensorValue operator behavior. Reproduce the uint8 example from the issue, then verify that the named ops and TensorValue's and, or, and xor provide the expected bitwise behavior for the torch MAX backend without the workaround described in issue 5299.

Written by the indexing model from the issue text.

Assessment

Tech stack
pytorch
Domain
backend, machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.