Adding bitwise ONNX ops

crates/burn-import/SUPPORTED-ONNX-OPS.md

@@ -25,11 +25,11 @@ functionality.
 | [AveragePool2d][12]              | ✅             | ✅           |
 | [BatchNormalization][14]         | ✅             | ✅           |
 | [Bernoulli][15]                  | ❌             | ✅           |
-| [BitShift][16]                   | ❌             | ✅           |
-| [BitwiseAnd][17]                 | ❌             | ✅           |
-| [BitwiseNot][18]                 | ❌             | ✅           |
-| [BitwiseOr][19]                  | ❌             | ✅           |
-| [BitwiseXor][20]                 | ❌             | ✅           |
+| [BitShift][16]                   | ✅             | ✅           |
+| [BitwiseAnd][17]                 | ✅             | ✅           |
+| [BitwiseNot][18]                 | ✅             | ✅           |
+| [BitwiseOr][19]                  | ✅             | ✅           |
+| [BitwiseXor][20]                 | ✅             | ✅           |
 | [BlackmanWindow][21]             | ❌             | ❌           |
 | [Cast][22]                       | ✅             | ✅           |
 | [CastLike][23]                   | ❌             | ❌           |
@@ -92,7 +92,7 @@ functionality.
 | [If][77]                         | ❌             | ✅           |
 | [Im][78]                         | ❌             | ❌           |
 | [InstanceNormalization][79]      | ✅             | ✅           |
-| [IsInf][80]                      | ❌             | ❌           |
+| [IsInf][80]                      | ❌             | ✅           |
 | [IsNaN][81]                      | ❌             | ✅           |
 | [LayerNormalization][82]         | ✅             | ✅           |
 | [LeakyRelu][83]                  | ✅             | ✅           |

crates/burn-import/onnx-tests/build.rs

@@ -14,6 +14,17 @@ fn main() {
         .input("tests/avg_pool1d/avg_pool1d.onnx")
         .input("tests/avg_pool2d/avg_pool2d.onnx")
         .input("tests/batch_norm/batch_norm.onnx")
+        .input("tests/bitshift/bitshift_left.onnx")
+        .input("tests/bitshift/bitshift_left_scalar.onnx")
+        .input("tests/bitshift/bitshift_right.onnx")
+        .input("tests/bitshift/bitshift_right_scalar.onnx")
+        .input("tests/bitwise_and/bitwise_and.onnx")
+        .input("tests/bitwise_and/bitwise_and_scalar.onnx")
+        .input("tests/bitwise_not/bitwise_not.onnx")
+        .input("tests/bitwise_or/bitwise_or.onnx")
+        .input("tests/bitwise_or/bitwise_or_scalar.onnx")
+        .input("tests/bitwise_xor/bitwise_xor.onnx")
+        .input("tests/bitwise_xor/bitwise_xor_scalar.onnx")
         .input("tests/cast/cast.onnx")
         .input("tests/ceil/ceil.onnx")
         .input("tests/clip/clip.onnx")

crates/burn-import/onnx-tests/tests/bitshift/bitshift.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+# used to generate model: onnx-tests/tests/bitshift/bitshift_left.onnx and bitshift_right.onnx
+
+import onnx
+
+def build_model(direction: str = "LEFT", scalar_shift: bool = False):
+    op_type = "BitShift"
+    direction_attr = "LEFT" if direction == "LEFT" else "RIGHT"
+
+    nodes = [
+        onnx.helper.make_node(
+            op_type,
+            inputs=["x", "shift"],
+            outputs=["output"],
+            name=f"/{op_type}",
+            direction=direction_attr
+        ),
+    ]
+
+    # Both tensor and scalar versions have the same input structure
+    # The scalar version will be handled by the burn runtime with scalar input
+    inputs = [
+        onnx.helper.make_value_info(
+            name="x",
+            type_proto=onnx.helper.make_tensor_type_proto(
+                elem_type=onnx.TensorProto.INT32, shape=[4]
+            ),
+        ),
+        onnx.helper.make_value_info(
+            name="shift",
+            type_proto=onnx.helper.make_tensor_type_proto(
+                elem_type=onnx.TensorProto.INT32, shape=[] if scalar_shift else [4]
+            ),
+        ),
+    ]
+
+    return onnx.helper.make_model(
+        ir_version=8,
+        opset_imports=[onnx.helper.make_operatorsetid("", 18)],
+        graph=onnx.helper.make_graph(
+            name="main_graph",
+            nodes=nodes,
+            inputs=inputs,
+            outputs=[
+                onnx.helper.make_value_info(
+                    name="output",
+                    type_proto=onnx.helper.make_tensor_type_proto(
+                        elem_type=onnx.TensorProto.INT32, shape=[4]
+                    ),
+                )
+            ],
+            initializer=[]
+        ),
+    )
+
+def export_bitshift(direction: str = "LEFT"):
+    # Regular tensor version
+    onnx_model = build_model(direction, scalar_shift=False)
+    file_name = f"bitshift_{direction.lower()}.onnx"
+
+    onnx.checker.check_model(onnx_model)
+    onnx.save(onnx_model, file_name)
+    print(f"Finished exporting model to {file_name}")
+
+    # Scalar version
+    onnx_model_scalar = build_model(direction, scalar_shift=True)
+    file_name_scalar = f"bitshift_{direction.lower()}_scalar.onnx"
+
+    onnx.checker.check_model(onnx_model_scalar)
+    onnx.save(onnx_model_scalar, file_name_scalar)
+    print(f"Finished exporting model to {file_name_scalar}")
+
+if __name__ == "__main__":
+    for direction in ["LEFT", "RIGHT"]:
+        export_bitshift(direction)

crates/burn-import/onnx-tests/tests/bitshift/mod.rs

@@ -0,0 +1,73 @@
+// Include the models for this node type
+use crate::include_models;
+include_models!(
+    bitshift_left,
+    bitshift_left_scalar,
+    bitshift_right,
+    bitshift_right_scalar
+);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use burn::tensor::{Int, Tensor, TensorData};
+
+    type Backend = burn_ndarray::NdArray<f32>;
+
+    #[test]
+    fn bitshift_left_tensors() {
+        // Initialize the model with weights (loaded from the exported file)
+        let device = Default::default();
+        let model: bitshift_left::Model<Backend> = bitshift_left::Model::new(&device);
+        // Run the model
+        let input1 = Tensor::<Backend, 1, Int>::from_ints([1, 2, 3, 4], &device);
+        let input2 = Tensor::<Backend, 1, Int>::from_ints([1, 1, 2, 2], &device);
+        let output = model.forward(input1, input2);
+        let expected = TensorData::from([2i64, 4, 12, 16]);
+
+        output.to_data().assert_eq(&expected, true);
+    }
+
+    #[test]
+    fn bitshift_left_scalar_tensor() {
+        // Initialize the model with weights (loaded from the exported file)
+        let device = Default::default();
+        let model: bitshift_left_scalar::Model<Backend> = bitshift_left_scalar::Model::new(&device);
+        // Run the model
+        let input1 = Tensor::<Backend, 1, Int>::from_ints([1, 2, 3, 4], &device);
+        let scalar = 2;
+        let output = model.forward(input1, scalar);
+        let expected = TensorData::from([4i64, 8, 12, 16]);
+
+        output.to_data().assert_eq(&expected, true);
+    }
+
+    #[test]
+    fn bitshift_right_tensors() {
+        let device = Default::default();
+        let model: bitshift_right::Model<Backend> = bitshift_right::Model::new(&device);
+
+        // Run the model
+        let input1 = Tensor::<Backend, 1, Int>::from_ints([1, 2, 3, 4], &device);
+        let input2 = Tensor::<Backend, 1, Int>::from_ints([1, 1, 2, 2], &device);
+        let output = model.forward(input1, input2);
+        let expected = TensorData::from([0i64, 1, 0, 1]);
+
+        output.to_data().assert_eq(&expected, true);
+    }
+
+    #[test]
+    fn bitshift_right_scalar_tensor() {
+        // Initialize the model with weights (loaded from the exported file)
+        let device = Default::default();
+        let model: bitshift_right_scalar::Model<Backend> =
+            bitshift_right_scalar::Model::new(&device);
+        // Run the model
+        let input1 = Tensor::<Backend, 1, Int>::from_ints([1, 2, 3, 4], &device);
+        let scalar = 2;
+        let output = model.forward(input1, scalar);
+        let expected = TensorData::from([0i64, 0, 0, 1]);
+
+        output.to_data().assert_eq(&expected, true);
+    }
+}

crates/burn-import/onnx-tests/tests/bitwise_and/bitwise_and.onnx

@@ -0,0 +1,18 @@
+pytorch2.5.1:m
+'
+
x
+
youtput/BitwiseAnd"
+BitwiseAnd
+main_graphZ
+
x
+
+
+Z
+
y
+
+
+b
+output
+
+
+B

crates/burn-import/onnx-tests/tests/bitwise_and/bitwise_and.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python3
+# used to generate model: onnx-tests/tests/bitwise_and/bitwise_and.onnx
+
+import torch
+import onnx
+
+def export_bitwise_and():
+    class BitwiseAndModel(torch.nn.Module):
+        def forward(self, x, y):
+            if isinstance(y, int):
+                # If y is a scalar, convert it to a tensor
+                y = torch.tensor([y], dtype=x.dtype)
+            return torch.bitwise_and(x, y)
+
+    model = BitwiseAndModel()
+    x = torch.tensor([1, 2, 3, 4], dtype=torch.int32)
+    y = torch.tensor([4, 3, 2, 1], dtype=torch.int32)
+    torch.onnx.export(
+        model,
+        (x, y),
+        "bitwise_and.onnx",
+        opset_version=18,
+        input_names=["x", "y"],
+        output_names=["output"],
+    )
+
+    # Scalar version
+    and_scalar = 2  # Scalar shift value
+    torch.onnx.export(
+        model,
+        (x, and_scalar),
+        f"bitwise_and_scalar.onnx",
+        opset_version=18,
+        input_names=["x", "y"],
+        output_names=["output"],
+    )
+
+if __name__ == "__main__":
+    export_bitwise_and()

crates/burn-import/onnx-tests/tests/bitwise_and/mod.rs

@@ -0,0 +1,35 @@
+// Include the models for this node type
+use crate::include_models;
+include_models!(bitwise_and, bitwise_and_scalar);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use burn::tensor::{Int, Tensor, TensorData};
+
+    type Backend = burn_ndarray::NdArray<f32>;
+
+    #[test]
+    fn bitwise_and_tensors() {
+        let device = Default::default();
+        let model: bitwise_and::Model<Backend> = bitwise_and::Model::new(&device);
+        // Run the model
+        let input1 = Tensor::<Backend, 1, Int>::from_ints([1, 2, 3, 4], &device);
+        let input2 = Tensor::<Backend, 1, Int>::from_ints([1, 1, 2, 2], &device);
+        let output = model.forward(input1, input2);
+        let expected = TensorData::from([1i64, 0, 2, 0]);
+        output.to_data().assert_eq(&expected, true);
+    }
+
+    #[test]
+    fn bitwise_and_scalar_tensor() {
+        let device = Default::default();
+        let model: bitwise_and_scalar::Model<Backend> = bitwise_and_scalar::Model::new(&device);
+        // Run the model
+        let input1 = Tensor::<Backend, 1, Int>::from_ints([1, 2, 3, 4], &device);
+        let scalar = 2;
+        let output = model.forward(input1, scalar);
+        let expected = TensorData::from([0i64, 2, 2, 0]);
+        output.to_data().assert_eq(&expected, true);
+    }
+}

crates/burn-import/onnx-tests/tests/bitwise_not/bitwise_not.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+
+# used to generate model: onnx-tests/tests/bitwise_not/bitwise_not.onnx
+
+import onnx
+
+def build_model():
+    return onnx.helper.make_model(
+        ir_version=8,
+        opset_imports=[onnx.helper.make_operatorsetid("", 18)],
+        graph=onnx.helper.make_graph(
+            name="main_graph",
+            nodes=[
+                onnx.helper.make_node(
+                    "BitwiseNot",
+                    inputs=["input"],
+                    outputs=["output"],
+                    name="/BitwiseNot"
+                ),
+            ],
+            inputs=[
+                onnx.helper.make_value_info(
+                    name="input",
+                    type_proto=onnx.helper.make_tensor_type_proto(
+                        elem_type=onnx.TensorProto.INT32, shape=[1, 4]
+                    ),
+                ),
+            ],
+            outputs=[
+                onnx.helper.make_value_info(
+                    name="output",
+                    type_proto=onnx.helper.make_tensor_type_proto(
+                        elem_type=onnx.TensorProto.INT32, shape=[1, 4]
+                    ),
+                )
+            ]
+        ),
+    )
+
+def main():
+    onnx_model = build_model()
+    file_name = "bitwise_not.onnx"
+
+    onnx.checker.check_model(onnx_model)  # Ensure valid ONNX
+    onnx.save(onnx_model, file_name)  # Save the model
+    print(f"Finished exporting model to {file_name}")
+
+if __name__ == "__main__":
+    main()

crates/burn-import/onnx-tests/tests/bitwise_not/mod.rs

@@ -0,0 +1,22 @@
+// Include the models for this node type
+use crate::include_models;
+include_models!(bitwise_not);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use burn::tensor::{Int, Tensor, TensorData};
+
+    type Backend = burn_ndarray::NdArray<f32>;
+
+    #[test]
+    fn bitwise_not_tensors() {
+        let device = Default::default();
+        let model: bitwise_not::Model<Backend> = bitwise_not::Model::new(&device);
+        // Run the model
+        let input = Tensor::<Backend, 2, Int>::from_ints([[1, 2, 3, 4]], &device);
+        let output = model.forward(input);
+        let expected = TensorData::from([[-2i64, -3, -4, -5]]);
+        output.to_data().assert_eq(&expected, true);
+    }
+}