# Copyright (c) 2023, Apple Inc. All rights reserved.
#
# Use of this source code is governed by a BSD-3-clause license that can be
# found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause
from coremltools.converters.mil.mil import types
from coremltools.converters.mil.mil.input_type import InputSpec, TensorInputType
from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op
from coremltools.converters.mil.mil.ops.defs.iOS15.linear import linear as _linear_iOS15
from coremltools.converters.mil.mil.ops.defs.iOS15.linear import matmul as _matmul_iOS15
from coremltools.converters.mil.mil.ops.defs.iOS17 import _IOS17_TARGET
[docs]@register_op(opset_version=_IOS17_TARGET)
class linear(_linear_iOS15):
"""
A version of ``linear`` for iOS 17+. The only difference between this version and the
iOS 15 :py:class:`~.iOS15.linear.linear` is that the ``weight`` and ``bias`` may have a
different dtype than the input/output.
Parameters
----------
x: tensor<[\*D, D_in], T> (Required)
* ``1 <= rank <= 3``.
* ``0 <= rank(*D) <= 2``.
weight: const tensor<[D_out, D_in], U> (Required)
bias: const tensor<[D_out], U> (Optional)
* Default to ``0``.
Returns
-------
tensor<[\*D, D_out], T>
* Same rank as the input ``x``.
Attributes
----------
T: fp16, fp32, i32
U: fp16, fp32, i32
"""
input_spec = InputSpec(
x=TensorInputType(type_domain="T"),
weight=TensorInputType(const=True, type_domain="U"),
bias=TensorInputType(const=True, optional=True, type_domain="U"),
)
type_domains = {
"T": (types.fp16, types.fp32, types.int32),
"U": (types.fp16, types.fp32, types.int32),
}
[docs]@register_op(opset_version=_IOS17_TARGET)
class matmul(_matmul_iOS15):
"""
A version of ``matmul`` for iOS 17+. The only difference between this version and the
iOS 15 :py:class:`~.iOS15.linear.matmul` is that the ``x`` and ``y`` can have a different
dtypes when one of them is const.
Parameters
----------
x: tensor<[\*, K1], T> (Required)
* ``x`` must be 1-D or higher.
y: tensor<[\*, K2], U> (Required)
* ``y`` must be 1-D or higher.
transpose_x: const bool (Optional)
* Default to ``False``.
* Use ``True`` to transpose the last two dimensions of ``x`` before multiplication.
It has no effect when ``x`` is 1-D.
transpose_y: const bool (Optional)
* Default to ``False``.
* Use ``True`` to transpose the last two dimensions of ``y`` before multiplication.
It has no effect when ``y`` is 1-D.
Returns
-------
tensor<\*, V>
* Scalar or tensor output.
* When ``x`` and ``y`` are both const or both non-const, it should follow ios15 behavior
that ``x``, ``y``, and ``output`` all have the same dtype.
When one of x and y is const, the output dtype should be the same as the non-const one.
Attributes
----------
T: fp16, fp32, i32
U: fp16, fp32, i32
"""
input_spec = InputSpec(
x=TensorInputType(type_domain="T"),
y=TensorInputType(type_domain="U"),
transpose_x=TensorInputType(const=True, optional=True, type_domain=types.bool),
transpose_y=TensorInputType(const=True, optional=True, type_domain=types.bool),
)
type_domains = {
"T": (types.fp16, types.fp32, types.int32),
"U": (types.fp16, types.fp32, types.int32),
}
def type_inference(self):
x_is_const = self.x.op is not None and self.x.op.op_type == "const"
y_is_const = self.y.op is not None and self.y.op.op_type == "const"
if x_is_const == y_is_const and self.x.dtype != self.y.dtype:
is_const_str = "const" if x_is_const else "non-const"
raise ValueError(
f'In op "matmul", when x and y are both {is_const_str}, their dtype '
f"need to match, but got x as {types.builtin_to_string(self.x.dtype)} "
f"and y as {types.builtin_to_string(self.y.dtype)}"
)
inferred_type = super().type_inference()
if x_is_const != y_is_const:
# The output dtype should be the same as the non-const one.
output_dtype = self.x.dtype if y_is_const else self.y.dtype
inferred_type = types.tensor(output_dtype, inferred_type.get_shape())
return inferred_type