trunc, fmod and Mod ONNX ops (tracel-ai#3767)

* Add trunc operation for float tensors

Introduces a new `trunc` method for float tensors, which truncates each element toward zero. Updates documentation and test suites to include and verify the new operation.

* Add fmod operation for float tensors

Introduces the fmod and fmod_scalar methods to compute floating-point remainders for float tensors. Updates documentation and adds comprehensive tests to verify correct behavior for various input scenarios.

* Add ONNX Mod operator support and tests

Implements support for the ONNX Mod operator, including both fmod (C-style) and remainder (Python-style) behaviors. Adds ModNode and related codegen, updates node registration, conversion, and rank inference logic. Introduces comprehensive ONNX test models and Rust test cases for tensor, scalar, and broadcasting scenarios.

* Handle IEEE 754 edge cases in fmod tensor ops

Improves the fmod and fmod_scalar tensor operations to correctly handle IEEE 754 special cases, including NaN, infinity, and zero divisors. Adds comprehensive tests for these edge cases and broadcasting behavior to ensure compliance and correctness.

* Fix formatting

* Mark ONNX Mod op as supported

Updated SUPPORTED-ONNX-OPS.md to indicate that the Mod operation is now supported for both import and export.

* Refactor fmod handling and test assertions

* Improve ONNX mod tests with ReferenceEvaluator checks

* Fix zero times infinity check in fmod implementation

* Improve trunc method to handle IEEE 754 special cases

Updated the trunc implementation to correctly handle special cases per IEEE 754, including preserving the sign of zero, infinities, and NaN. Added corresponding tests to verify correct behavior for these cases.

Author: antimora

burn-book/src/building-blocks/tensor.md

@@ -275,6 +275,8 @@ Those operations are only available for `Float` tensors.
 | `tensor.erf()`                               | `tensor.erf()`                             |
 | `tensor.exp()`                               | `tensor.exp()`                             |
 | `tensor.floor()`                             | `tensor.floor()`                           |
+| `tensor.fmod(other)`                         | `tensor.fmod(other)`                       |
+| `tensor.fmod_scalar(scalar)`                 | `tensor.fmod(scalar)`                      |
 | `tensor.from_floats(floats, device)`         | N/A                                        |
 | `tensor.from_full_precision(tensor)`         | N/A                                        |
 | `tensor.int()`                               | Similar to `tensor.to(torch.long)`         |
@@ -295,6 +297,7 @@ Those operations are only available for `Float` tensors.
 | `tensor.tan()`                               | `tensor.tan()`                             |
 | `tensor.tanh()`                              | `tensor.tanh()`                            |
 | `tensor.to_full_precision()`                 | `tensor.to(torch.float)`                   |
+| `tensor.trunc()`                             | `tensor.trunc()`                           |
 | `tensor.var(dim)`                            | `tensor.var(dim)`                          |
 | `tensor.var_bias(dim)`                       | N/A                                        |
 | `tensor.var_mean(dim)`                       | N/A                                        |

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

@@ -118,7 +118,7 @@ functionality.
 | [MelWeightMatrix][103]           | ❌             | ❌           |
 | [Min][104]                       | ✅             | ✅           |
 | [Mish][105]                      | ❌             | ❌           |
-| [Mod][106]                       | ❌             | ❌           |
+| [Mod][106]                       | ✅             | ✅           |
 | [Mul][107]                       | ✅             | ✅           |
 | [Multinomial][108]               | ❌             | ❌           |
 | [Neg][109]                       | ✅             | ✅           |

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

@@ -107,6 +107,12 @@ fn main() {
         .input("tests/div/div_shape.onnx")
         .input("tests/div/div_shape_tensor.onnx")
         .input("tests/div/div_broadcast.onnx")
+        .input("tests/mod/modulo.onnx")
+        .input("tests/mod/mod_scalar.onnx")
+        .input("tests/mod/mod_remainder.onnx")
+        .input("tests/mod/mod_fmod.onnx")
+        .input("tests/mod/mod_broadcast_fixed.onnx")
+        .input("tests/mod/mod_broadcast_remainder_fixed.onnx")
         .input("tests/dropout/dropout.onnx")
         .input("tests/equal/equal.onnx")
         .input("tests/equal/equal_shape.onnx")

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

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+
+# used to generate model: onnx-tests/tests/mod/mod.onnx
+
+import torch
+import torch.nn as nn
+
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, x, y):
+        # Compute the modulo operation
+        z = torch.fmod(x, y)
+        return z
+
+
+def main():
+    # Export to onnx
+    model = Model()
+    model.eval()
+    device = torch.device("cpu")
+    onnx_name = "mod.onnx"
+
+    # Create dummy inputs with proper shapes
+    dummy_x = torch.randn(2, 3, 4, device=device)
+    dummy_y = torch.randn(2, 3, 4, device=device)
+
+    torch.onnx.export(model, (dummy_x, dummy_y), onnx_name,
+                      verbose=False, opset_version=16)
+
+    print("Finished exporting model to {}".format(onnx_name))
+
+    # Output some test data for use in the test
+    test_x = torch.tensor([[[[5.3, -5.3, 7.5, -7.5]]]])
+    test_y = torch.tensor([[[[2.0, 2.0, 3.0, 3.0]]]])
+
+    print("Test input x: {}".format(test_x))
+    print("Test input y: {}".format(test_y))
+    output = model.forward(test_x, test_y)
+    print("Test output: {}".format(output))
+
+
+if __name__ == '__main__':
+    main()

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

@@ -0,0 +1,170 @@
+// Import the shared macro
+use crate::include_models;
+include_models!(
+    modulo,
+    mod_scalar,
+    mod_remainder,
+    mod_fmod,
+    mod_broadcast_fixed,
+    mod_broadcast_remainder_fixed
+);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use burn::tensor::{Tensor, TensorData};
+
+    use crate::backend::TestBackend;
+
+    #[test]
+    fn mod_tensor_by_tensor() {
+        // Initialize the model
+        let device = Default::default();
+        let model: modulo::Model<TestBackend> = modulo::Model::new(&device);
+
+        // Run the model
+        let input_x = Tensor::<TestBackend, 3>::from_floats([[[5.3, -5.3, 7.5, -7.5]]], &device);
+        let input_y = Tensor::<TestBackend, 3>::from_floats([[[2.0, 2.0, 3.0, 3.0]]], &device);
+        let output = model.forward(input_x, input_y);
+
+        // Expected output: fmod(x, y) for each element
+        // Using the actual computed values from Python
+        let expected = TensorData::from([[[1.3000002f32, -1.3000002, 1.5, -1.5]]]);
+
+        output.to_data().assert_eq(&expected, true);
+    }
+
+    #[test]
+    fn mod_tensor_by_scalar() {
+        // Initialize the model
+        let device = Default::default();
+        let model: mod_scalar::Model<TestBackend> = mod_scalar::Model::new(&device);
+
+        // Run the model
+        let input_x = Tensor::<TestBackend, 4>::from_floats([[[[5.3, -5.3, 7.5, -7.5]]]], &device);
+        let scalar = 2.0f64;
+        let output = model.forward(input_x, scalar);
+
+        // Expected output: fmod(x, 2.0) for each element
+        // Using the actual computed values from Python
+        let expected = TensorData::from([[[[1.3000002f32, -1.3000002, 1.5, -1.5]]]]);
+
+        output.to_data().assert_eq(&expected, true);
+    }
+
+    #[test]
+    fn mod_remainder() {
+        // Test fmod=0 (Python-style remainder)
+        let device = Default::default();
+        let model: mod_remainder::Model<TestBackend> = mod_remainder::Model::new(&device);
+
+        let input_x = Tensor::<TestBackend, 3>::from_floats([[[5.3, -5.3, 7.5, -7.5]]], &device);
+        let input_y = Tensor::<TestBackend, 3>::from_floats([[[2.0, 2.0, 3.0, 3.0]]], &device);
+        let output = model.forward(input_x, input_y);
+
+        // Expected: Python-style remainder where sign follows divisor
+        // remainder(5.3, 2.0) = 1.3, remainder(-5.3, 2.0) = 0.7
+        // remainder(7.5, 3.0) = 1.5, remainder(-7.5, 3.0) = 1.5
+        let expected = TensorData::from([[[1.3000002f32, 0.6999998, 1.5, 1.5]]]);
+        output.to_data().assert_eq(&expected, true);
+    }
+
+    #[test]
+    fn mod_fmod() {
+        // Test fmod=1 (C-style fmod)
+        let device = Default::default();
+        let model: mod_fmod::Model<TestBackend> = mod_fmod::Model::new(&device);
+
+        let input_x = Tensor::<TestBackend, 3>::from_floats([[[5.3, -5.3, 7.5, -7.5]]], &device);
+        let input_y = Tensor::<TestBackend, 3>::from_floats([[[2.0, 2.0, 3.0, 3.0]]], &device);
+        let output = model.forward(input_x, input_y);
+
+        // Expected: fmod operation where sign follows dividend
+        let expected = TensorData::from([[[1.3000002f32, -1.3000002, 1.5, -1.5]]]);
+        output.to_data().assert_eq(&expected, true);
+    }
+
+    #[test]
+    fn mod_broadcast() {
+        // Test broadcasting with fmod=1
+        let device = Default::default();
+        let model: mod_broadcast_fixed::Model<TestBackend> =
+            mod_broadcast_fixed::Model::new(&device);
+
+        let input_x = Tensor::<TestBackend, 2>::from_floats(
+            [
+                [5.0, -7.0, 8.0, -9.0],
+                [4.0, -6.0, 10.0, -11.0],
+                [3.0, -5.0, 12.0, -13.0],
+            ],
+            &device,
+        );
+
+        let input_y = Tensor::<TestBackend, 4>::from_floats(
+            [
+                [[
+                    [3.0, 3.0, 3.0, 3.0],
+                    [3.0, 3.0, 3.0, 3.0],
+                    [3.0, 3.0, 3.0, 3.0],
+                ]],
+                [[
+                    [4.0, 4.0, 4.0, 4.0],
+                    [4.0, 4.0, 4.0, 4.0],
+                    [4.0, 4.0, 4.0, 4.0],
+                ]],
+            ],
+            &device,
+        );
+
+        let output = model.forward(input_x, input_y);
+
+        // Check shape and sample values
+        assert_eq!(output.dims(), [2, 1, 3, 4]);
+
+        // Check first batch, first row
+        let data = output.to_data();
+        let values = data.as_slice::<f32>().unwrap();
+        // fmod(5.0, 3.0) = 2.0, fmod(-7.0, 3.0) = -1.0, etc.
+        assert!((values[0] - 2.0).abs() < 0.001);
+        assert!((values[1] - (-1.0)).abs() < 0.001);
+        assert!((values[2] - 2.0).abs() < 0.001);
+        assert!((values[3] - 0.0).abs() < 0.001);
+    }
+
+    #[test]
+    fn mod_broadcast_remainder() {
+        // Test broadcasting with fmod=0 (remainder)
+        let device = Default::default();
+        let model: mod_broadcast_remainder_fixed::Model<TestBackend> =
+            mod_broadcast_remainder_fixed::Model::new(&device);
+
+        let input_x =
+            Tensor::<TestBackend, 3>::from_floats([[[7.5], [-8.5], [9.5], [-10.5]]], &device);
+
+        let input_y = Tensor::<TestBackend, 3>::from_floats(
+            [
+                [[3.0, 4.0, -3.0, -4.0, 5.0]],
+                [[3.0, 4.0, -3.0, -4.0, 5.0]],
+                [[3.0, 4.0, -3.0, -4.0, 5.0]],
+            ],
+            &device,
+        );
+
+        let output = model.forward(input_x, input_y);
+
+        // Check shape
+        assert_eq!(output.dims(), [3, 4, 5]);
+
+        // Check first row values for Python-style remainder
+        // remainder(7.5, 3.0) = 1.5, remainder(7.5, 4.0) = 3.5
+        // remainder(7.5, -3.0) = -1.5, remainder(7.5, -4.0) = -0.5
+        // remainder(7.5, 5.0) = 2.5
+        let data = output.to_data();
+        let values = data.as_slice::<f32>().unwrap();
+        assert!((values[0] - 1.5).abs() < 0.001);
+        assert!((values[1] - 3.5).abs() < 0.001);
+        assert!((values[2] - (-1.5)).abs() < 0.001);
+        assert!((values[3] - (-0.5)).abs() < 0.001);
+        assert!((values[4] - 2.5).abs() < 0.001);
+    }
+}

crates/burn-import/onnx-tests/tests/mod/mod_broadcast.onnx

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+
+# used to generate model: onnx-tests/tests/mod/mod_broadcast.onnx
+# Tests broadcasting with fmod=1
+
+import numpy as np
+import onnx
+from onnx import helper, TensorProto
+
+def main():
+    # Create ONNX model with Mod operator using broadcasting
+    # Different rank tensors: 2D and 4D
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4])  # 2D
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [2, 1, 3, 4])  # 4D
+
+    # Output tensor will have the broadcasted shape
+    z = helper.make_tensor_value_info('z', TensorProto.FLOAT, [2, 1, 3, 4])
+
+    # Create Mod node with fmod=1 (C-style fmod)
+    mod_node = helper.make_node(
+        'Mod',
+        inputs=['x', 'y'],
+        outputs=['z'],
+        fmod=1  # C-style fmod
+    )
+
+    # Create the graph
+    graph_def = helper.make_graph(
+        [mod_node],
+        'mod_broadcast_model',
+        [x, y],
+        [z],
+    )
+
+    # Create the model with opset version 16
+    model_def = helper.make_model(
+        graph_def,
+        producer_name='onnx-tests',
+        opset_imports=[helper.make_operatorsetid("", 16)]
+    )
+
+    # Save the model
+    onnx_name = "mod_broadcast.onnx"
+    onnx.save(model_def, onnx_name)
+    onnx.checker.check_model(onnx_name)
+    print(f"Finished exporting model to {onnx_name}")
+
+    # Test with onnx.reference.ReferenceEvaluator
+    try:
+        from onnx.reference import ReferenceEvaluator
+
+        # Create test data
+        test_x = np.array([[5.0, -7.0, 8.0, -9.0],
+                           [4.0, -6.0, 10.0, -11.0],
+                           [3.0, -5.0, 12.0, -13.0]]).astype(np.float32)
+
+        test_y = np.array([[[[3.0, 3.0, 3.0, 3.0],
+                             [3.0, 3.0, 3.0, 3.0],
+                             [3.0, 3.0, 3.0, 3.0]]],
+                           [[[4.0, 4.0, 4.0, 4.0],
+                             [4.0, 4.0, 4.0, 4.0],
+                             [4.0, 4.0, 4.0, 4.0]]]]).astype(np.float32)
+
+        # Run inference with ReferenceEvaluator
+        sess = ReferenceEvaluator(model_def)
+        result = sess.run(None, {"x": test_x, "y": test_y})
+
+        print(f"Test input x shape: {test_x.shape}")
+        print(f"Test input y shape: {test_y.shape}")
+        print(f"Result shape: {result[0].shape}")
+        print(f"Sample output values (result[0,0,0,:]): {result[0][0, 0, 0, :]}")
+
+        # Verify expected results for fmod operation
+        test_x_broadcast = np.broadcast_to(test_x, test_y.shape)
+        expected_result = np.fmod(test_x_broadcast, test_y)
+        np.testing.assert_allclose(result[0], expected_result, rtol=1e-5)
+        print("Test passed: Results match expected fmod values")
+
+    except ImportError:
+        print("onnx.reference not available, skipping inference test")
+
+if __name__ == '__main__':
+    main()

crates/burn-import/onnx-tests/tests/mod/mod_broadcast.py

@@ -0,0 +1,22 @@
+pytorch2.8.0:�

+5
+onnx::Mod_0
+onnx::Mod_1
2/Mod"Mod*
+fmod
�

+main_graphZ
+onnx::Mod_0
+

+
+Z%
+onnx::Mod_1
+

+
+

+
+b
+
2
+

+
+

+
+B