feat: Mod

onnx2torch/node_converters/__init__.py

@@ -26,6 +26,7 @@
 from onnx2torch.node_converters.max_pool import *
 from onnx2torch.node_converters.mean import *
 from onnx2torch.node_converters.min_max import *
+from onnx2torch.node_converters.mod import *
 from onnx2torch.node_converters.neg import *
 from onnx2torch.node_converters.nms import *
 from onnx2torch.node_converters.pad import *

onnx2torch/node_converters/mod.py

@@ -0,0 +1,33 @@
+__all__ = [
+    'OnnxMod',
+]
+
+import torch
+from torch import nn
+from onnx2torch.node_converters.registry import add_converter
+from onnx2torch.onnx_graph import OnnxGraph
+from onnx2torch.onnx_node import OnnxNode
+from onnx2torch.utils.common import OnnxToTorchModule, OperationConverterResult, onnx_mapping_from_node
+
+
+class OnnxMod(nn.Module, OnnxToTorchModule):  # pylint: disable=missing-class-docstring
+    def __init__(self, fmod: int):
+        super().__init__()
+        self.fmod = fmod
+
+        if self.fmod not in [0, 1]:
+            raise ValueError(f'OnnxMod fom must be 0 or 1, but get {self.fmod}')
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:  # pylint: disable=missing-function-docstring
+        return torch.fmod(x, y) if self.fmod else torch.remainder(x, y)
+
+
+@add_converter(operation_type='Mod', version=10)
+@add_converter(operation_type='Mod', version=13)
+def _(node: OnnxNode, graph: OnnxGraph) -> OperationConverterResult:  # pylint: disable=unused-argument
+    node_attributes = node.attributes
+    fmod = node_attributes.get('fmod', 0)
+    return OperationConverterResult(
+        torch_module=OnnxMod(fmod=fmod),
+        onnx_mapping=onnx_mapping_from_node(node=node),
+    )

operators.md

@@ -79,7 +79,7 @@ Minimal tested opset version 9, maximum tested opset version 16, recommended ops
 | MaxUnpool                 | N         |                                                                                                                                                                     |
 | Mean                      | Y         |                                                                                                                                                                     |
 | Min                       | Y         |                                                                                                                                                                     |
-| Mod                       | N         |                                                                                                                                                                     |
+| Mod                       | Y         |                                                                                                                                                                     |
 | Mul                       | Y         |                                                                                                                                                                     |
 | Multinomial               | N         |                                                                                                                                                                     |
 | Neg                       | Y         |                                                                                                                                                                     |

tests/node_converters/mod_test.py

@@ -0,0 +1,40 @@
+from typing import List
+
+import numpy as np
+import onnx
+import pytest
+from tests.utils.common import check_onnx_model, make_model_from_nodes
+
+
+@pytest.mark.parametrize(
+    'dividend',
+    [
+        [-4, 7, 5, 4, -7, 8],
+        [-4.3, 7.2, 5.0, 4.3, -7.2, 8.0],
+    ],
+)
+@pytest.mark.parametrize(
+    'divisor',
+    [
+        [2, -3, 8, -2, 3, 5],
+        [2.1, -3.4, 8.0, -2.1, 3.4, 5.0],
+    ],
+)
+@pytest.mark.parametrize('fmod', [0, 1])
+def test_mod(  # pylint: disable=missing-function-docstring
+    dividend: List[float],
+    divisor: List[float],
+    fmod: int,
+) -> None:
+    x_variants = np.array(dividend).astype(np.float32 if fmod else np.int32)
+    y_variants = np.array(divisor).astype(np.float32 if fmod else np.int32)
+
+    test_inputs = {'x': x_variants, 'y': y_variants}
+
+    node = onnx.helper.make_node(op_type='Mod', inputs=['x', 'y'], outputs=['z'], fmod=fmod)
+    model = make_model_from_nodes(
+        nodes=node,
+        initializers={},
+        inputs_example=test_inputs,
+    )
+    check_onnx_model(model, test_inputs)