Good First Issue: Add MLX Op Handler for `aten.bitwise_xor`

[metascroy]

🚀 The feature, motivation and pitch

Good First Issue: Add MLX Op Handler for aten.bitwise_xor

Summary

Add support for aten.bitwise_xor in the MLX delegate. This op performs element-wise bitwise XOR and is needed for bit manipulation, checksums, encryption primitives, and toggle operations.

Background

The MLX delegate currently has no handler for bitwise_xor. MLX has native support via mlx::core::bitwise_xor which works for both boolean and integer tensors.

Approach: New schema node + runtime

Add a BitwiseXorNode to handle this via MLX's native bitwise_xor.

Steps

1. Add node to backends/mlx/serialization/schema.fbs

   table BitwiseXorNode {
     a: Tid;
     b: Tid;
     out: Tid;
   }

   Add BitwiseXorNode to the OpNode union (append only, do not reorder).

2. Regenerate serialization code

   python backends/mlx/serialization/generate.py

3. Add C++ runtime exec function in backends/mlx/runtime/MLXInterpreter.h

   inline void exec_bitwise_xor(
       const BitwiseXorNode& n, ExecutionState& st, StreamOrDevice s) {
     auto a = st.get_tensor(n.a());
     auto b = st.get_tensor(n.b());
     auto out = mlx::core::bitwise_xor(a, b, s);
     st.set_tensor(n.out(), out);
   }

4. Add handler in backends/mlx/ops.py

   Add to the _BINARY_OPS table:

   # In _BINARY_OPS list, add:
   ([torch.ops.aten.bitwise_xor.Tensor, torch.ops.aten.bitwise_xor.Scalar], BitwiseXorNode, "aten.bitwise_xor", True),

5. Add test in backends/mlx/test/test_ops.py

   Use the _BINARY_OP_TESTS table:

   # Add to _BINARY_OP_TESTS list:
   {"op_name": "bitwise_xor_bool", "op_fn": torch.bitwise_xor, 
    "shapes": _SHAPES_3, "dtypes": [torch.bool], 
    "input_fn_a": _bool_input_fn(), "input_fn_b": _bool_input_fn()},
   {"op_name": "bitwise_xor_int", "op_fn": torch.bitwise_xor, 
    "shapes": _SHAPES_3, "dtypes": [torch.int32, torch.int64], 
    "input_fn_a": _int_input_fn(0, 256), "input_fn_b": _int_input_fn(0, 256)},

Running tests

python -m executorch.backends.mlx.test.run_all_tests -k bitwise_xor

References

• MLX C++: array bitwise_xor(const array &a, const array &b, StreamOrDevice s = {})
• PyTorch signature: bitwise_xor(Tensor self, Tensor other) -> Tensor
• Supported dtypes: int8, int16, int32, int64, uint8, bool
• Note: For bool tensors, bitwise_xor is equivalent to logical_xor

Alternatives

No response

Additional context

No response

RFC (Optional)

No response

[yadferhad]

I’d like to work on this issue. I plan to add a BitwiseXorNode path for aten.bitwise_xor in the MLX delegate, wire it through the generated schema/runtime loader, and add bool and integer regression tests in the MLX op test suite.

[metascroy]

I’d like to work on this issue. I plan to add a BitwiseXorNode path for aten.bitwise_xor in the MLX delegate, wire it through the generated schema/runtime loader, and add bool and integer regression tests in the MLX op test suite.

Great! I see @nanookclaw has put up a PR for this already. If that isn't you, happy to assign one of the similar issues to you (e.g., support for bitwise_and) if you're interested