Skip to content

Commit 6aa864e

Browse files
authored
Add ability to add MIL ops to ET that are not yet in coremltools (#12648)
This introduces torch_ops.py in ET's CoreML compiler. This lets us add critical MIL ops that are not yet in ET's version of coremltools.
1 parent 0fbd6d4 commit 6aa864e

File tree

2 files changed

+28
-0
lines changed

2 files changed

+28
-0
lines changed

backends/apple/coreml/compiler/coreml_preprocess.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@
2828
logger = logging.getLogger(__name__)
2929
logger.setLevel(logging.WARNING)
3030

31+
from executorch.backends.apple.coreml.compiler.torch_ops import * # noqa: F401, F403
32+
3133

3234
class COMPILE_SPEC_KEYS(Enum):
3335
COMPUTE_UNITS = "compute_units"
Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
# This file registers torch ops that are not yet in coremltools, or are in a more recent version of
8+
# coremltools than is used by ExecuTorch. Each op registered here should have a link to a PR in coremltools that adds
9+
# the op to the coremltools library.
10+
11+
from coremltools.converters.mil.frontend.torch.ops import transpose, unbind
12+
from coremltools.converters.mil.frontend.torch.torch_op_registry import (
13+
register_torch_op,
14+
)
15+
16+
17+
# https://github.com/apple/coremltools/pull/2556
18+
@register_torch_op(override=False)
19+
def transpose_copy(context, node):
20+
transpose(context, node)
21+
22+
23+
# https://github.com/apple/coremltools/pull/2557
24+
@register_torch_op(override=False)
25+
def unbind_copy(context, node):
26+
unbind(context, node)

0 commit comments

Comments
 (0)