Adding support for torch.minimum and torch.maximum (#1231)

Co-authored-by: Côme Weber <[email protected]>
Co-authored-by: Aseem Wadhwa <[email protected]>

coremltools/converters/mil/frontend/torch/ops.py

@@ -742,6 +742,26 @@ def max_pool3d(context, node):
     _max_pool(context, node, inputs)
 
 
+@register_torch_op
+def minimum(context, node):
+    inputs = _get_inputs(context, node, expected=2)
+    assert len(node.outputs) == 1
+    x = context[node.inputs[0]]
+    y = context[node.inputs[1]]
+    out = mb.minimum(x=x, y=y, name=node.name)
+    context.add(out)
+
+
+@register_torch_op
+def maximum(context, node):
+    inputs = _get_inputs(context, node, expected=2)
+    assert len(node.outputs) == 1
+    x = context[node.inputs[0]]
+    y = context[node.inputs[1]]
+    out = mb.maximum(x=x, y=y, name=node.name)
+    context.add(out)
+
+
 @register_torch_op
 def div(context, node):
     inputs = _get_inputs(context, node, expected=2)

coremltools/converters/mil/frontend/torch/test/test_torch_ops.py

@@ -966,6 +966,29 @@ def test_max_pool3d(
         )
         self.run_compare_torch(input_shape, model, backend=backend)
 
+class TestMaximumMinimum(TorchBaseTest):
+
+    @pytest.mark.parametrize(
+        "input_shape, mode, backend",
+        itertools.product(
+            [(2, 5, 7, 3), (3, 2, 9)],
+            ["minimum", "maximum"],
+            backends,
+        ),
+    )
+    def test_minimum_maximum(self, input_shape, mode, backend):
+        class TestModel(torch.nn.Module):
+            def forward(self, x, y):
+                if mode == "minimum":
+                    return torch.minimum(x, y)
+                elif mode == "maximum":
+                    return torch.maximum(x, y)
+                else:
+                    raise ValueError(f"Unsupported mode: {mode}")
+
+        model = TestModel()
+        self.run_compare_torch([input_shape] * 2, model, backend=backend)
+
 class TestPoolSymbolicInput(TorchBaseTest):
     def test_max_pool(self):
         model = nn.MaxPool2d(