Skip to content

Core ML: run ops on the CPU when Core ML cannot take or return their values #24317

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shoumikhin wants to merge 3 commits into
base: main
Choose a base branch
from
+1,598 −42
Open

Core ML: run ops on the CPU when Core ML cannot take or return their values #24317

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
aaeb50b
Move Core ML ops to the CPU at values Core ML cannot take or return
shoumikhin Nov 3, 2026
85c5f90
Stop keeping an op whole when decomposing makes a call Core ML rejects
shoumikhin Nov 3, 2026
d892975
Use fill_like for Core ML fills shaped like a tensor
shoumikhin Nov 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 17 additions & 4 deletions backends/apple/coreml/compiler/torch_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@

import numpy as np
import torch as _torch
from coremltools import _logger
from coremltools import _logger, target
from coremltools.converters.mil import Builder as mb
from coremltools.converters.mil.frontend import _utils
from coremltools.converters.mil.frontend.torch import ops as _ct_ops
Expand All @@ -31,6 +31,9 @@
register_torch_op,
)
from coremltools.converters.mil.mil import types
from coremltools.converters.mil.mil.block import (
is_current_opset_version_compatible_with,
)
from coremltools.converters.mil.mil.types.symbolic import is_symbolic
from executorch.exir.dim_order_utils import get_memory_format

Expand Down Expand Up @@ -92,6 +95,16 @@ def _add_fill(context, node, shape, value):
context.add(_make_fill_op(shape, _core_ml_type(node)(value), node.name))


def _add_fill_like(context, node, x, value):
# A fill of a shape known at conversion becomes a constant as big as x in
# the model, which fill_like avoids from iOS 16.
if is_current_opset_version_compatible_with(target.iOS16):
value = _core_ml_type(node)(value)
context.add(mb.fill_like(ref_tensor=x, value=value, name=node.name))
else:
_add_fill(context, node, mb.shape(x=x), value)


# coremltools gives these fills the type of their fill value or of their input,
# or makes them float, instead of the dtype torch gives them, so an int or bool
# fill could come out as float. https://github.com/apple/coremltools/pull/2865
Expand Down Expand Up @@ -119,7 +132,7 @@ def _fill_like(context, node, convert):
inputs = _get_inputs(context, node, min_expected=2)
if _torch_dtype(node) is None or inputs[1].val is None:
return convert(context, node)
_add_fill(context, node, mb.shape(x=inputs[0]), inputs[1].val)
_add_fill_like(context, node, inputs[0], inputs[1].val)


@register_torch_op(override=True)
Expand Down Expand Up @@ -153,15 +166,15 @@ def zeros_like(context, node):
if _torch_dtype(node) is None:
return _ct_ops.zeros_like(context, node)
x = _get_inputs(context, node, min_expected=1)[0]
_add_fill(context, node, mb.shape(x=x), 0)
_add_fill_like(context, node, x, 0)


@register_torch_op(override=True)
def ones_like(context, node):
if _torch_dtype(node) is None:
return _ct_ops.ones_like(context, node)
x = _get_inputs(context, node, min_expected=1)[0]
_add_fill(context, node, mb.shape(x=x), 1)
_add_fill_like(context, node, x, 1)


@register_torch_op(override=True)
Expand Down
Loading
Loading