diff --git a/coremltools/converters/mil/frontend/torch/ops.py b/coremltools/converters/mil/frontend/torch/ops.py index d2470423c..3e8a6694d 100644 --- a/coremltools/converters/mil/frontend/torch/ops.py +++ b/coremltools/converters/mil/frontend/torch/ops.py @@ -5796,7 +5796,11 @@ def bitwise_and(context, node): # which torch.export emits for `tensor | tensor` / `tensor ^ tensor`. These are # common when building boolean attention masks (e.g. Gemma combines a causal # mask with a padding mask via __or__). -@register_torch_op(torch_alias=["or"]) +# +# "ior" is the post-sanitize form of "aten::__ior__", the in-place `|=`. +# `sanitize_op_kind` only strips a trailing `_`, so the leading `i` in the +# `__i__` family is preserved -- the alias has to be listed explicitly. +@register_torch_op(torch_alias=["or", "ior"]) def bitwise_or(context, node): _bitwise_as_logical_if_boolean(context, node, "bitwise_or", logical_or) diff --git a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py index 00dd7f5f9..4b7bb38c8 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py +++ b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py @@ -13754,6 +13754,41 @@ def forward(self, x, y): input_as_shape=False, ) + @pytest.mark.parametrize( + "compute_unit, backend", + itertools.product(compute_units, backends), + ) + def test_ior_operator(self, compute_unit, backend): + # Regression test for issue #2584: TorchScript trace of `z |= y` + # records `aten::__ior__`, which sanitizes to "ior" and must be + # registered as an alias of bitwise_or. Previously gemma-3-1b-it + # conversion failed with + # "PyTorch convert function for op '__ior__' not implemented." + # + # Notes on scope: + # * Core ML inputs are immutable, so the model clones an input + # before mutating it; without the clone, `run_compare_torch` + # would raise the unrelated user-input-mutation guard. + # * torch.export decomposes `__ior__` into `clone + bitwise_or`, + # so the new alias is only reachable via the TorchScript path. + class TestModel(torch.nn.Module): + def forward(self, x, y): + z = x.clone() + z |= y + return z + + input_shape = (2, 3) + input_data_x = torch.rand(*input_shape) > 0.2 + input_data_y = torch.rand(*input_shape) < 0.8 + self.run_compare_torch( + [input_data_x, input_data_y], + TestModel(), + frontend=TorchFrontend.TORCHSCRIPT, + backend=backend, + compute_unit=compute_unit, + input_as_shape=False, + ) + class TestBitwiseXor(TorchBaseTest): @pytest.mark.parametrize(