class scatter_nd(Operation):
"""
Scatter ``updates`` to ``data`` at locations ``indices``.
The ``indices`` is a K-dim tensor, where ``indices[i_0,...,i_{K-2}]`` defines a
slice of ``data``, ``K = rank(indices)``, and ``data[indices[i_0, ..., i_{K-2}]]``
has rank ``rank(data) - indices.shape[-1]``.
* Example: ``mode == update``: The ``output`` is set to ``data`` initially, and
the op updates ``output`` as follows:
.. math::
output[indices[i_0, ..., i_{K-2}]]= updates[indices[i_0, ..., i_{K-2}]]
* Example: ``mode == add``. The update rule is:
.. math::
output[indices[i_0, ..., i_{K-2}]] += updates[indices[i_0, ..., i_{K-2}]]
Parameters
----------
data: tensor<\*D,T> (Required)
indices: tensor<\*K,i32> (Required)
updates: tensor<\*K, T> (Required)
* Must be the shape as ``K[:-1]+data.shape[K[-1]:]``.
mode: const string (Optional)
* Default to ``add``.
* Can be the following modes: ``update``, ``add``, ``sub``, ``mul``,
``div``, ``max``, ``min``.
Returns
-------
tensor<\*D,T>
* A tensor with the same shape and type as ``data``.
Attributes
----------
T: fp16, fp32, i32
"""
input_spec = InputSpec(
data=TensorInputType(),
indices=IntTensorInputType(),
updates=TensorInputType(),
mode=StringInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(mode="add", )
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
assert self.indices.shape[-1] <= self.data.rank
expected_updates_shape = (self.indices.shape[:-1] +
self.data.shape[self.indices.shape[-1]:])
assert is_compatible_symbolic_vector(self.updates.shape,
tuple(expected_updates_shape))
return self.data.sym_type
class prelu(Operation):
"""
Where ``i = 1 ... C``, if ``x_i > 0``, return ``x_i`` , otherwise return ``alpha_i * x_i``.
Parameters
----------
x: tensor<[B, C, 1..3], T> (Required)
* x must have rank 4 or rank 3 or rank 5, i.e. a shape of (B,C,H) or (B,C,H,W) or (B,C,D,H,W)
alpha: const tensor<[C], T>, (Required)
* The length of alpha must match the second dimension of x (channel dimension)
Returns
-------
tensor<[B, C, 1..3], T>
* A tensor of the same shape as ``x``.
Attributes
----------
T: fp32, fp16
"""
input_spec = InputSpec(
x=TensorInputType(),
alpha=TensorInputType(const=True),
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
@precondition(allow=VALUE)
def value_inference(self):
alpha_br = self.alpha.val
for i in range(1, len(self.x.shape)):
alpha_br = np.expand_dims(alpha_br, i)
x_pos = np.maximum(self.x.val, 0)
b = np.minimum(self.x.val, 0)
return x_pos + b * alpha_br
def type_inference(self):
if self.x.rank not in (3, 4, 5):
raise ValueError(
"prelu op: x must be rank 3 or 4 or 5, instead it is of rank {}"
.format(len(self.x.shape)))
if len(self.alpha.val.shape) != 1:
raise ValueError("alpha should be rank 1")
if self.x.shape[1] != self.alpha.val.shape[0]:
raise ValueError(
"Size of dimension 1 of alpha should be the same as " +
"the size of dimension 1 of x.")
if self.x.rank in (3, 5):
# check whether all alpha values are the same or not
are_values_same = np.where(
np.abs(self.alpha.val - self.alpha.val[0]) > 1e-5)[0].size == 0
if not are_values_same:
raise ValueError(
"prelu op: rank 3 or rank 5 input is only supported when all the values of alpha are same,"
"which is not the case here")
return self.x.sym_type
class softplus_parametric(Operation):
"""
Return ``alpha_i * log( 1 + e^( beta_i * x_i ) )``, where ``i = 1 ... C``.
Parameters
----------
x: tensor<[b, C, n, m], T> (Required)
alpha: const tensor<[C], fp32> (Required)
beta: const tensor<[C], fp32> (Required)
Returns
-------
tensor<[b, C, n, m], T>
* A tensor of the same shape as ``x``.
Attributes
----------
T: fp32
"""
input_spec = InputSpec(
x=TensorInputType(),
alpha=TensorInputType(const=True),
beta=TensorInputType(const=True),
)
def __init__(self, **kwargs):
super(softplus_parametric, self).__init__(**kwargs)
@precondition(allow=VALUE)
def value_inference(self):
alpha_br = np.copy(self.alpha.val)
beta_br = np.copy(self.beta.val)
for i in range(1, len(self.x.val.shape)):
alpha_br = np.expand_dims(alpha_br, i)
beta_br = np.expand_dims(beta_br, i)
return alpha_br * np.log(1 + np.exp(self.x.val * beta_br))
def type_inference(self):
if len(self.x.shape) < 3:
raise ValueError("x should be at least rank 3")
if len(self.alpha.val.shape) != 1:
raise ValueError("alpha should be rank 1")
if self.x.shape[1] != self.alpha.val.shape[0]:
raise ValueError(
"Size of dimension 0 of alpha should be the same as "
+ "the size of dimension 1 of x."
)
if len(self.beta.val.shape) != 1:
raise ValueError("beta should be rank 1")
if self.x.shape[1] != self.beta.val.shape[0]:
raise ValueError(
"Size of dimension 0 of beta should be the same as "
+ "the size of dimension 1 of x."
)
return self.x.sym_type
class tf_lstm_block_cell(TfLSTMBase):
input_spec = (
InputSpec(x=TensorInputType(), ) +
TfLSTMBase.input_spec # [batch, input_dim]
)
def __init__(self, **kwargs):
super(tf_lstm_block_cell, self).__init__(**kwargs)
def type_inference(self):
self._check_peephole_weights()
# all return shapes are [batch, hidden_dim]
ret_shape = self.c_prev.shape
dtype = self.x.dtype
# See
# https://www.tensorflow.org/api_docs/python/tf/raw_ops/LSTMBlockCell
# All returned shapes are [batch, hidden_dim]
return (
types.tensor(dtype, ret_shape), # i
types.tensor(dtype, ret_shape), # cs
types.tensor(dtype, ret_shape), # f
types.tensor(dtype, ret_shape), # o
types.tensor(dtype, ret_shape), # ci
types.tensor(dtype, ret_shape), # co
types.tensor(dtype, ret_shape),
) # h
class tf_make_list(Operation):
input_spec = InputSpec(
init_length=IntInputType(optional=True),
dynamic_length=BoolInputType(optional=True),
elem_shape=TensorInputType(const=True, optional=True),
dtype=StringInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
init_length=1,
dynamic_length=True,
dtype="fp32",
)
def __init__(self, **kwargs):
super(tf_make_list, self).__init__(**kwargs)
def type_inference(self):
init_length = self.init_length.val
if self.elem_shape is None or self.elem_shape.sym_val is None:
return types.list(
types.unknown,
init_length=init_length,
dynamic_length=self.dynamic_length.val,
)
builtin_dtype = types.string_to_builtin(self.dtype.val)
if builtin_dtype is None:
raise ValueError("Unsupported dtype {}".format(self.dtype.val))
elem_type = types.tensor(builtin_dtype, self.elem_shape.sym_val)
return types.list(elem_type,
init_length=init_length,
dynamic_length=self.dynamic_length.val)
class non_zero(Operation):
"""
Returns the indices of the elements in the given tensor that are non-zero.
Parameters
----------
x: tensor<\*?, T> (Required)
* Tensor, values selected at indices where its values is not equal to ``0``.
Returns
-------
tensor<[N, R], int32>
* 2-dimensional tensor contains indices of elements that are non-zero.
Each row is the index for a non-zero value.
* ``N`` is the number of non-zero elements, ``R`` is the rank of the input.
Attributes
----------
T: fp32, int32
"""
input_spec = InputSpec(x=TensorInputType())
def __init__(self, **kwargs):
super(non_zero, self).__init__(**kwargs)
def type_inference(self):
shape = tuple([get_new_symbol(), self.x.rank])
return types.tensor(types.int32, shape)
@precondition(allow=VALUE)
def value_inference(self):
return np.transpose(np.nonzero(self.x.val))
class space_to_depth(Operation):
"""
Rearrange elements in a tensor from spatial into depth (channel) dimension.
Parameters
----------
x: tensor<[n, C, H, W], T> (Required)
* Input tensor of rank ``4``.
block_size: const<i32> (Required)
* The size of the spatial block. Must be greater than ``1`` and divisible
by spatial dimensions ``H, W``.
Returns
-------
tensor<[n, C x block_size^2, H / block_size, W / block_size], T>
* Where ``b`` is the block size.
Attributes
----------
T: fp32
"""
input_spec = InputSpec(
x=TensorInputType(),
block_size=IntInputType(const=True),)
def __init__(self, **kwargs):
super(space_to_depth, self).__init__(**kwargs)
def type_inference(self):
x_type = self.x.dtype
n, c, h, w = self.x.shape
bs = self.block_size.val
ret_shape = (n, c * (bs * bs), h // bs, w // bs)
return types.tensor(x_type, ret_shape)
class tf_lstm_block(TfLSTMBase):
"""
Apply LSTM to an input sequence
"""
input_spec = (
InputSpec(
seq_len=IntInputType(), # int
x=TensorInputType(), # [padded_len, batch, input_dim]
) + TfLSTMBase.input_spec)
def __init__(self, **kwargs):
super(tf_lstm_block, self).__init__(**kwargs)
def type_inference(self):
self._check_peephole_weights()
padded_len = self.x.shape[0]
ret_shape = [padded_len] + list(self.c_prev.shape)
dtype = self.x.dtype
# All returned shapes are [padded_len, b, hidden_dim]
return (
types.tensor(dtype, ret_shape), # i
types.tensor(dtype, ret_shape), # cs
types.tensor(dtype, ret_shape), # f
types.tensor(dtype, ret_shape), # o
types.tensor(dtype, ret_shape), # ci
types.tensor(dtype, ret_shape), # co
types.tensor(dtype, ret_shape),
) # h
class list_scatter(Operation):
"""
Scatter ``values`` to ``ls`` at locations ``indices``.
Parameters
----------
ls: List[*] (Required)
indices: tensor<num_updates, i32> (Required)
* Indices of ``ls`` to scatter to.
* Elements of ``indices`` must be in ``[0, ls.length)`` at runtime.
* If indices are greater than or equal to the list length, the list is
dynamically resized.
value: <*,T> (Optional)
* Element value to write, which must match the element shape of ``ls``.
* Default is ``None``.
Returns
-------
List[*]
* Updated list.
Attributes
----------
T: fp32, i32, bool
"""
input_spec = InputSpec(
ls=ListInputType(),
indices=IntTensorInputType(),
value=TensorInputType(),
)
def __init__(self, **kwargs):
super(list_scatter, self).__init__(**kwargs)
def type_inference(self):
num_indices = self.indices.shape[0]
num_values = self.value.shape[0]
if num_values != num_indices:
raise ValueError("Cannot scatter {} values to {} indices".format(
num_values, num_indices))
list_elem_type = self.ls.elem_type
value_type = self.value.sym_type
dynamic_length = self.ls.dynamic_length
init_length = self.ls.init_length
elem_type = types.tensor(value_type.get_primitive(),
value_type.get_shape()[1:])
if list_elem_type == types.unknown:
# fill in the elem type using value's type info.
return types.list(elem_type,
dynamic_length=dynamic_length,
init_length=init_length)
if not types.is_subtype(elem_type, list_elem_type):
msg = "Elem type mismatch: ls elem type {} vs " + "value type {}"
raise ValueError(msg.format(list_elem_type, elem_type))
return self.ls.sym_type
class prelu(Operation):
"""
Where ``i = 1 ... C``, if ``x_i > 0``, return ``x_i`` , otherwise return ``alpha_i * x_i``.
Parameters
----------
x: tensor<[b, C, n, m], T> (Required)
alpha: const tensor<[C], T>, (Required)
Returns
-------
tensor<[b, C, n, m], fp32>
* A tensor of the same shape as ``x``.
Attributes
----------
T: fp32
"""
input_spec = InputSpec(
x=TensorInputType(),
alpha=TensorInputType(const=True),
)
def __init__(self, **kwargs):
super(prelu, self).__init__(**kwargs)
@precondition(allow=VALUE)
def value_inference(self):
alpha_br = self.alpha.val
for i in range(1, len(self.x.shape)):
alpha_br = np.expand_dims(alpha_br, i)
x_pos = np.maximum(self.x.val, 0)
b = np.minimum(self.x.val, 0)
return x_pos + b * alpha_br
def type_inference(self):
if len(self.x.shape) < 3:
raise ValueError("x should be at least rank 3")
if len(self.alpha.val.shape) != 1:
raise ValueError("alpha should be rank 1")
if self.x.shape[1] != self.alpha.val.shape[0]:
raise ValueError(
"Size of dimension 1 of alpha should be the same as " +
"the size of dimension 1 of x.")
return self.x.sym_type
class reverse_sequence(Operation):
"""
Reverse variable length slices for specified axes / dimensions of the input
tensor. This op first slices input tensor along the ``batch_axis`` dimension, then
partially reverses the elements along the ``seq_axis`` for the first ``lengths[i]``
elements.
Parameters
----------
x: tensor<\*?, T> (Required)
* Input tensor.
lengths: tensor<L, i32> (Required)
* 1-dimensional tensor of length ``x.shape[batch_axis]`` specifying the length
of the sequence to reverse.
* Values must be in range ``[0, x.shape[seq_axis]]``.
seq_axis: const<i32> (Optional)
* The dimension to reverse.
* Defaults to ``0``.
batch_axis: const<i32> (Optional)
* Dimension for slicing.
* Defaults to ``0``.
Returns
-------
tensor<\*?, T>
* Same type and shape as the input tensor.
Attributes
----------
T: fp32
References
----------
`tf.reverse_sequence <https://www.tensorflow.org/api_docs/python/tf/reverse_sequence>`_
"""
input_spec = InputSpec(
x=TensorInputType(),
lengths=IntTensorInputType(),
seq_axis=IntInputType(const=True, optional=True),
batch_axis=IntInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
seq_axis=0,
batch_axis=0)
def __init__(self, **kwargs):
super(reverse_sequence, self).__init__(**kwargs)
def type_inference(self):
return self.x.sym_type
@precondition(allow=VALUE)
def value_inference(self):
raise NotImplementedError("TODO")
class flatten2d(Operation):
"""
Flattens input tensor into 2d tensor by flattening dimensions before and
after the provided axis.
Parameters
----------
x: tensor<[*d], T> (Required)
* Input tensor.
* axis: const<f32> (Optional)
* Defaults to ``1``.
* Negative axis is supported.
Returns
-------
tensor<d_prior, d_post, T>
* ``d_prior`` is product of dimensions ``x[:axis]``
* ``d_post`` is product of dimensions ``x[axis:]``
Examples
--------
1. ``input_shape = (3, ), axis = -1, output_shape = (1, 3)``
2. ``input_shape = (3, ), axis = 1, output_shape = (3, 1)``
3. ``input_shape = (4, 3), axis = -1, output_shape = (4, 3)``
4. ``input_shape = (2, 3, 2), axis = -1, output_shape = (6, 2)``
5. ``input_shape = (5, 5, 2), axis = 1, output_shape = (5, 10)``
Attributes
----------
T: fp16, fp32, i32, bool
"""
input_spec = InputSpec(x=TensorInputType(),
axis=IntInputType(const=True, optional=True))
def default_inputs(self):
return DefaultInputs(axis=1, )
def __init__(self, **kwargs):
super(flatten2d, self).__init__(**kwargs)
def type_inference(self):
shape = list(self.x.shape)
axis = self.axis.val
dim_pre_axis = np.prod(shape[:axis])
dim_post_axis = np.prod(shape[axis:])
new_shape = [dim_pre_axis, dim_post_axis]
return types.tensor(self.x.dtype, tuple(new_shape))
@precondition(allow=VALUE | SYMBOL)
def value_inference(self):
shape = self.x.shape
axis = self.axis.val
dim_pre_axis = np.prod(shape[:axis])
dim_post_axis = np.prod(shape[axis:])
return self.x.val.reshape(dim_pre_axis, dim_post_axis)
class classify(Operation):
"""
Presence of this op indicates that the model is of type classifier,
which accordingly constructs the model output, that is, the predicted class label
and the output probability dictionary. The parameters of this op are set
based on the attributes set for the class "coremltools.ClassifierConfig" by the user.
The outputs of this op cannot be used by another op.
Parameters
----------
* probabilities: tensor<[* , ProbT]> (Required)
* a tensor in the graph, which is used to compute the classifier output(s)
* This is the tensor whose values are mapped to the class labels and used for
* constructing the predicted class label and the output dictionary of class names
* and values
* classes: list<*, ClassT> (Required)
* list of classes
Returns
-------
* <classT>
* Dict[classT, probT]
Attributes
----------
ProbT: fp32
ClassT: int64, str
"""
input_spec = InputSpec(
probabilities=TensorInputType(),
classes=ListInputType(const=True),
)
def __init__(self, **kwargs):
super(classify, self).__init__(**kwargs)
def type_inference(self):
# check the type of "classes"
if not types.is_list(self.classes.sym_type):
msg = "'classes' in the op 'classify' must be of type list. Instead it is {}."
raise ValueError(msg.format(self.classes.sym_type.__type_info__()))
classes_elem_type = self.classes.elem_type
if classes_elem_type not in {types.str, types.int64}:
msg = "Type of elements in 'classes' in the op 'classify' must be either str or int64. Instead it is {}."
raise ValueError(msg.format(classes_elem_type.__type_info__()))
# check that the size of "classes" is compatible with the size of "probabilities"
if not any_symbolic(self.probabilities.shape):
size = np.prod(self.probabilities.shape)
if len(self.classes.val) != size:
msg = "In op 'classify', number of classes must match the size of the tensor corresponding to 'probabilities'."
raise ValueError(msg)
return classes_elem_type, types.dict(classes_elem_type, types.double)
class list_write(Operation):
"""
Write a value into index ``index`` of ``ls``.
Parameters
----------
ls: List (Required)
index: <i32> (Required)
* Size of the list.
value: <*,T> (Optional)
* Element value to write, which must match the element shape of ``ls``.
* Default is ``None``.
Returns
-------
List[*]
Attributes
----------
T: fp32, i32, bool
"""
input_spec = InputSpec(
ls=ListInputType(),
index=IntInputType(),
value=TensorInputType(),
)
def __init__(self, **kwargs):
super(list_write, self).__init__(**kwargs)
def type_inference(self):
list_elem_type = self.ls.elem_type
value_type = self.value.sym_type
dynamic_length = self.ls.dynamic_length
init_length = self.ls.init_length
if list_elem_type is None:
# fill in the elem type using value's type info.
return types.list(value_type,
init_length=init_length,
dynamic_length=dynamic_length)
if list_elem_type == types.unknown:
msg = "Input ls elem type unknown. Override with {}"
logging.warning(msg.format(value_type))
return types.list(value_type,
init_length=init_length,
dynamic_length=dynamic_length)
if not types.is_subtype(value_type, list_elem_type):
msg = "Elem type mismatch: ls elem type {} vs " + "value type {}"
raise ValueError(
msg.format(list_elem_type.__type_info__(),
value_type.__type_info__()))
return self.ls.sym_type
class Pooling(Operation):
"""
Pooling Op Superclass
"""
input_spec = InputSpec(
x=TensorInputType(),
kernel_sizes=IntTensorInputType(const=True),
strides=IntTensorInputType(const=True, optional=True),
pad_type=StringInputType(const=True),
pad=IntTensorInputType(const=True, optional=True),
ceil_mode=BoolInputType(const=True, optional=True),
)
def default_inputs(self):
num_spatial_dims = self.x.rank - 2
return DefaultInputs(
strides=[1]*num_spatial_dims,
pad=[0]*2*num_spatial_dims,
ceil_mode=False,
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
ksize = self.kernel_sizes.val
x_shape = self.x.shape
D_in_rank = len(x_shape) - 2
strides = [1] * D_in_rank if self.strides is None else self.strides.val
pad_type = "valid" if self.pad_type is None else self.pad_type.val.lower()
if pad_type not in ["valid", "same", "custom"]:
raise ValueError("Unrecognized value of pad_type : {}".format(pad_type))
pad = None if self.pad is None else self.pad.val
D_in = x_shape[2:] # spatial dimensions
if self.ceil_mode.val:
if D_in_rank > 2:
raise ValueError('pool: ceil_mode only supported for 1D or 2D pool')
if pad_type == "same" and self.ceil_mode.val:
raise ValueError("ceil_mode must be False when pad_type==same")
if pad is not None:
for i in range(D_in_rank):
if pad[2*i] != pad[2*i+1]:
raise ValueError("Padding must be symmetric if ceil_mode is True")
D_out_shape = spatial_dimensions_out_shape(
pad_type=pad_type,
input_shape=D_in,
kernel_shape=ksize,
strides=strides,
custom_pad=pad,
ceil_mode=self.ceil_mode.val,
)
ret_shape = list(x_shape[:2]) + D_out_shape
return types.tensor(self.x.dtype, tuple(ret_shape))
class custom_torch_sparse_matmul(Operation):
# Defining input spec for current op
input_spec = InputSpec(
x=TensorInputType(),
y=TensorInputType(),
transpose_x=BoolInputType(const=True, optional=True),
transpose_y=BoolInputType(const=True, optional=True),
x_is_sparse=BoolInputType(const=True, optional=True),
y_is_sparse=BoolInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
transpose_x=False,
transpose_y=False,
x_is_sparse=False,
y_is_sparse=False,
)
# Specifying binding for custom op for specifying inputs,
# parameters required for creating custom op to be synced with Swift API
bindings = {
"class_name":
"SparseMatMul",
"input_order": ["x", "y"],
"parameters":
["transpose_x", "transpose_y", "x_is_sparse", "y_is_sparse"],
"description":
"Custom Sparse MatMul Layer",
}
def __init__(self, **kwargs):
super(TestCustomOp.custom_torch_sparse_matmul,
self).__init__(**kwargs)
def type_inference(self):
x_type = self.x.dtype
x_shape = self.x.shape
y_shape = self.y.shape
# For illustration purpose, assumming getting valid shape
# Ideally, should consider transpose_?, ?_is_sparse parameters into consideration
# for computing output shape
return types.tensor(x_type, [x_shape[0], y_shape[1]])
class random_categorical(Operation):
"""
Returns random values from a categorical distribution.
Parameters
----------
shape: <\*D_in, T>
* N-dimensional tensor, one of ``logits`` (event log-probabilities) or ``probs``
(event probabilities). The first ``N - 1`` dimensions specifies distributions,
and the last dimension represents a vector of probabilities.
mode: const<str> (Optional)
One of ``['logits', 'probs']``. Defaults to ``logits``.
size: const<i32> (Optional)
Number of samples to draw. Defaults to ``1``.
seed: const<i32> (Optional)
Seed to create a reproducible sequence of values across multiple invokes.
Returns
-------
<\*D_in[:-1] + [size], T>
* A tensor of the given target output shape filled with random values.
Attributes
----------
T: fp16, fp32
See Also
--------
random_bernoulli, random_normal, random_uniform
"""
input_spec = InputSpec(
x=TensorInputType(),
mode=StringInputType(const=True, optional=True),
size=IntInputType(const=True, optional=True),
seed=IntInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
mode="logits",
size=1,
seed=-1,
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
self.out_dtype = self.x.dtype
output_shape = self.x.shape[:-1] + (self.size.val,)
return types.tensor(self.out_dtype, output_shape)
class torch_upsample_bilinear(Operation):
"""
Upsample the spatial dimensions (last two dimensions) of the input by
scale factors using bilinear interpolation.
It corresponds to `torch.nn.functional.interpolate` function with `mode=bilinear`,
`recompute_scale_factor=True`, and input with flexible shape.
source: https://pytorch.org/docs/stable/_modules/torch/nn/functional.html#interpolate
Parameters
----------
x: tensor<[\*D, H1, W1],T> (Required)
* Must be rank ``3``.
output_height: i32
* Output height for the height dimension.
output_width: i32
* Output width for the width dimension.
aligh_corners: const<bool>
* The `aligh_corners` parameter for the original torch op.
Returns
-------
tensor<[\*D, H2, W2],T>
* Tensor with same type as the input.
* ``H2`` = output_height
* ``W2`` = output_width
Attributes
----------
T: fp32
"""
input_spec = InputSpec(
x=TensorInputType(),
output_height=IntOrFloatInputType(),
output_width=IntOrFloatInputType(),
align_corners=BoolInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
align_corners=True,
)
def __init__(self, **kwargs):
super(torch_upsample_bilinear, self).__init__(**kwargs)
def type_inference(self):
if self.x.rank < 3:
raise ValueError(
'input to the "torch_upsample_bilinear" op must have rank at least 3'
)
ret_shape = list(self.x.shape)
ret_shape[-1] = get_new_symbol()
ret_shape[-2] = get_new_symbol()
return types.tensor(self.x.dtype, ret_shape)
class sliding_windows(Operation):
"""
Return a tensor containing all windows of ``size``, separated by stride along the
given ``axis``.
Parameters
----------
x: tensor<[\*d0, d_axis, *dn], T>
* Input tensor.
axis: const<i32>
* Axis to perform the operation.
size: const<i32>
* Number of elements in the sliding window.
stride: const<i32> Optional
* Default to ``1``.
* The stride of the input elements in the sliding window.
Returns
-------
tensor<[\*d0, d_axis - size // stride + 1, size, \*dn], T>
* The output will be a tensor of rank ``N+1`` where ``N`` is the input tensor
rank.
Attributes
----------
T: fp16, fp32, i32
"""
input_spec = InputSpec(
x=TensorInputType(),
axis=IntInputType(const=True),
size=IntInputType(const=True),
stride=IntInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(stride=1)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
x_shape = self.x.shape
axis = self.axis.val
size = self.size.val
stride = self.stride.val
ret_shape = list(x_shape)
ret_shape[axis] = (x_shape[axis] - size) // stride + 1
pos_axis = axis if axis >= 0 else axis + self.x.rank
ret_shape.insert(pos_axis + 1, size)
return types.tensor(self.x.dtype, tuple(ret_shape))
class ReductionAxes(Operation):
"""
Reduction Op Superclasses
"""
input_spec = InputSpec(
x=TensorInputType(),
axes=IntTensorInputType(const=True, optional=True),
keep_dims=BoolInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
axes=None,
keep_dims=False,
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
x_type = self.x.dtype
x_shape = self.x.shape
axes = self.axes.val if self.axes is not None else None
if axes is None:
axes = range(self.x.rank)
keep_dims = self.keep_dims.val
reduced_shape = list(x_shape)
if keep_dims:
for i in axes:
reduced_shape[i] = 1
else:
# sort reverse so we can delete shape elements back to front
axes = [
axis if axis >= 0 else axis + len(reduced_shape)
for axis in axes
]
for i in sorted(axes)[::-1]:
reduced_shape.pop(i)
if len(reduced_shape) == 0:
return x_type # scalar
return types.tensor(x_type, tuple(reduced_shape))
@precondition(allow=VALUE)
def value_inference(self):
axes = tuple(self.axes.val) if self.axes is not None else None
return self.get_operator()(self.x.val,
axis=axes,
keepdims=self.keep_dims.val)
def get_operator(self):
raise NotImplementedError()
class transpose(Operation):
"""
Permute tensor ``x`` dimensions according to ``perm``.
Parameters
----------
x: tensor<\*?, T> (Required)
* Must be at least 1-D. ``x`` may have a symbolic shape.
perm: const<[rank(x)], i32> (Required)
* Permutation order. -rank(x) <= perm[I] < rank(x) for all perm entries.
Returns
-------
tensor<\*?,T>
* Tensor with same rank and type as ``x``.
Attributes
----------
T: fp16, fp32, i32, bool
References
----------
`torch.Tensor.permute <https://pytorch.org/docs/stable/tensors.html#torch.Tensor.permute>`_
"""
input_spec = InputSpec(
x=TensorInputType(),
perm=IntTensorInputType(const=True),
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
x_type = self.x.dtype
perm = self.perm.val
x_shape = np.array(self.x.shape)
if len(perm) != self.x.rank:
msg = "perm should have the same length as rank(x): {} != {}"
raise ValueError(msg.format(len(perm), self.x.rank))
if self.x.rank == 0:
return self.x.sym_type # scalar cannot be transposed
if any_variadic(self.x.shape):
ret_shape = get_new_variadic_symbol()
else:
ret_shape = x_shape[perm]
return types.tensor(x_type, tuple(ret_shape))
@precondition(allow=VALUE | SYMBOL)
def value_inference(self):
return np.transpose(self.x.val, axes=self.perm.val)
class reverse(Operation):
"""
Reverse the order of the input tensor ``x`` along specified ``axes`` (dimensions).
Parameters
----------
x: tensor<\*?, T> (Required)
* Input tensor.
axes: const<D, i32> (Optional)
* Dimension(s) to reverse. Each axis must be in the range ``[-rank(x), rank(x))``.
* Defaults to None (reverse on all dimensions).
Returns
-------
tensor<\*?, T>
* Same type and shape as the input tensor.
Attributes
----------
T: fp32
References
----------
See `tf.reverse <https://www.tensorflow.org/api_docs/python/tf/reverse>`_
and `TORCH <https://pytorch.org/docs/stable/torch.html#torch.flip>`_.
"""
input_spec = InputSpec(
x=TensorInputType(),
axes=IntTensorInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
axes=None,
)
def __init__(self, **kwargs):
super(reverse, self).__init__(**kwargs)
def type_inference(self):
return self.x.sym_type
@precondition(allow=VALUE)
def value_inference(self):
res = self.x.val
axes = self.axes.val if self.axes is not None else range(self.x.rank)
for axis in axes:
res = np.flip(res, axis=axis)
return res
class argsort(Operation):
"""
Returns a tensor containing the indices of the sorted values along a given axis
of the input tensor.
Parameters
----------
x: <\*?, T> (Required)
* Input tensor.
* axis: const<i32> (Optional)
* Default to ``-1`` (the last dimension).
* Axis to perform the operation.
* ascending: const<bool> (Optional)
* Default to ``False``, sort in descending order. ``True`` to sort in
ascending order.
Returns
-------
tensor<\*?, int32>
* Tensor containing the indices of the sorted values
Attributes
----------
T: fp32
"""
input_spec = InputSpec(
x=TensorInputType(),
axis=IntInputType(const=True, optional=True),
ascending=BoolInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
axis=-1,
ascending=False,
)
def __init__(self, **kwargs):
super(argsort, self).__init__(**kwargs)
def type_inference(self):
return types.tensor(types.int32, self.x.shape)
@precondition(allow=VALUE)
def value_inference(self):
# The default np argsort mode is ascending, which is opposite to MIL's argsort op.
if self.ascending.val:
return np.argsort(self.x.val, axis=self.axis.val)
return np.argsort(-self.x.val, axis=self.axis.val)
class torch_upsample_nearest_neighbor(Operation):
"""
Upsample the spatial dimensions (last two dimensions) of the input by
scale factors using nearest-neighbor interpolation.
It corresponds to `torch.nn.functional.interpolate` function with `mode=nearest`,
`recompute_scale_factor=True`, and input with flexible shape.
source: https://pytorch.org/docs/stable/_modules/torch/nn/functional.html#interpolate
Parameters
----------
x: tensor<[b, C, H1, W1],T> (Required)
* Must be rank ``4``.
output_height: i32
* Output height for the height dimension.
output_width: i32
* Output width for the width dimension.
Returns
-------
tensor<[b, C, H2, W2],T>
* Tensor with same type as the input.
* ``H2`` = output_height
* ``W2`` = output_width
Attributes
----------
T: fp32
"""
input_spec = InputSpec(
x=TensorInputType(),
output_height=IntOrFloatInputType(),
output_width=IntOrFloatInputType(),
)
def __init__(self, **kwargs):
super(torch_upsample_nearest_neighbor, self).__init__(**kwargs)
def type_inference(self):
if self.x.rank != 4:
raise ValueError(
'input to the "torch_upsample_nearest_neighbor" op must have rank 4'
)
ret_shape = list(self.x.shape)
ret_shape[2] = get_new_symbol()
ret_shape[3] = get_new_symbol()
return types.tensor(self.x.dtype, ret_shape)
class band_part(Operation):
"""
Returns a tensor setting everything outside a center band to zeros for the innermost
matrix. Special cases:
- ``band_part(x, 0, -1)`` returns upper triangular part.
- ``band_part(x, -1, 0)`` returns lower triangular part.
- ``band_part(x, 0, 0)`` returns diagonal.
Parameters
----------
x: tensor<\*?, T> (Required)
* Input tensor.
lower: const<i32> (Optional)
* Number of lower / below sub-diagonals to keep. If negative, keep entire
lower triangle.
* Defaults to ``-1`` (keep the entire lower triangle).
upper: const<i32> (Optional)
* Number of upper / above sub-diagonals to keep. If negative, keep entire
lower triangle.
* Defaults to ``-1`` (keep the entire upper triangle).
Returns
-------
tensor<\*?, T>
* Same type and shape as the input tensor.
Attributes
----------
T: fp16, fp32, i32, bool
"""
input_spec = InputSpec(
x=TensorInputType(),
lower=IntInputType(const=True, optional=True),
upper=IntInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(lower=-1, upper=-1)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
return self.x.sym_type
class gather_nd(Operation):
"""
Gather slices from ``x`` according to ``indices``, similar to `tf.gather_nd <https://www.tensorflow.org/api_docs/python/tf/gather_nd>`_.
The ``indices`` is a K-dim tensor, where ``indices[i_0,...,i_{K-2}]`` defines a slice
of ``x``:
.. math::
output[i_0, ..., i_{K-2}]= x[indices[i_0, ..., i_{K-2}]]
Where ``K = rank(indices)`` and ``x[indices[i_0, ..., i_{K-2}]]`` has rank
``rank(x) - indices.shape[-1]``.
Parameters
----------
x: tensor<\*D,T> (Required)
indices: tensor<\*K,i32> (Required)
Returns
-------
tensor<\*V,T>
* ``V = K[:-1] + D[K[-1]:]``, where ``D = x.shape`` and ``K = indices.shape``.
Attributes
----------
U: fp16, fp32, i32
References
----------
See `tf.gather_nd <https://www.tensorflow.org/api_docs/python/tf/gather_nd>`_.
"""
input_spec = InputSpec(
x=TensorInputType(),
indices=IntTensorInputType(),
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
assert self.indices.shape[-1] <= self.x.rank
out_type = self.x.dtype
out_shape = self.indices.shape[:-1] + self.x.shape[self.indices.
shape[-1]:]
return types.tensor(out_type, out_shape)
class ReductionAxis(Operation):
input_spec = InputSpec(
x=TensorInputType(),
axis=IntInputType(const=True, optional=True),
keep_dims=BoolInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
axis=-1,
keep_dims=False,
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def _find_reduced_shape(self):
x_shape = self.x.shape
axis = self.axis.val
reduced_shape = list(x_shape)
axis = axis if axis >= 0 else axis + len(reduced_shape)
if self.keep_dims.val:
reduced_shape[axis] = 1
else:
reduced_shape.pop(axis)
return reduced_shape
def type_inference(self):
x_type = self.x.dtype
reduced_shape = self._find_reduced_shape_and_axis()
return types.tensor(x_type, tuple(reduced_shape))
@precondition(allow=VALUE)
def value_inference(self):
tmp = self.get_operator()(self.x.val, axis=self.axis.val)
reduced_shape = self._find_reduced_shape()
if self.keep_dims.val:
tmp = np.reshape(tmp, reduced_shape)
return tmp
def get_operator(self):
raise NotImplementedError()
class resample(_resample_iOS15):
"""
iOS16 version of resample supports float16 coordinates
"""
input_spec = InputSpec(
x=TensorInputType(),
coordinates=ScalarOrTensorInputType(type_domain=(np.int32, np.float32, np.float16)),
sampling_mode=StringInputType(const=True),
padding_mode=StringInputType(const=True),
padding_value=FloatInputType(const=True),
coordinates_mode=StringInputType(const=True),
align_corners=BoolInputType(const=True),
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
return super().type_inference()
class pixel_unshuffle(Operation):
"""
Rearrange elements in a tensor from spatial dimensions into depth (channel).
It is basically the inverse operation of pixel_shuffle.
Equivalent to PyTorch's ``PixelUnshuffle``.
Parameters
----------
x: tensor<[n, C, H / f , W / f], T> (Required)
* Input tensor of rank ``4``.
downscale_factor: const<i32>
* Factor to decrease spatial resolution by.
Returns
-------
tensor<[n, C * f^2, H, W], T>
* Where ``f`` is the downscale factor.
Attributes
----------
T: fp32
References
----------
`torch.nn.PixelUnshuffle <https://pytorch.org/docs/stable/generated/torch.nn.PixelUnshuffle.html>`_
"""
input_spec = InputSpec(
x=TensorInputType(),
downscale_factor=ScalarOrTensorInputType(const=True,
type_domain=(np.uint32, )),
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
x_type = self.x.dtype
n, c, h, w = self.x.shape
f = self.downscale_factor.val
ret_shape = (n, c * f * f, h / f, w / f)
return types.tensor(x_type, ret_shape)
class tf_lstm_block_cell(TfLSTMBase):
"""
xh = [x, h_prev]
[i, ci, f, o] = xh * w + b
f = f + forget_bias
if not use_peephole:
wci = wcf = wco = 0
i = sigmoid(cs_prev .* wci + i)
f = sigmoid(cs_prev .* wcf + f)
ci = tanh(ci)
cs = ci .* i + cs_prev .* f
cs = clip(cs, cell_clip)
o = sigmoid(cs * wco + o)
co = tanh(cs)
h = co .* o
"""
input_spec = (
InputSpec(x=TensorInputType(), ) +
TfLSTMBase.input_spec # [batch, input_dim]
)
def __init__(self, **kwargs):
super(tf_lstm_block_cell, self).__init__(**kwargs)
def type_inference(self):
self._check_peephole_weights()
# all return shapes are [batch, hidden_dim]
ret_shape = self.c_prev.shape
dtype = self.x.dtype
# See
# https://www.tensorflow.org/api_docs/python/tf/raw_ops/LSTMBlockCell
# All returned shapes are [batch, hidden_dim]
return (
types.tensor(dtype, ret_shape), # i
types.tensor(dtype, ret_shape), # cs
types.tensor(dtype, ret_shape), # f
types.tensor(dtype, ret_shape), # o
types.tensor(dtype, ret_shape), # ci
types.tensor(dtype, ret_shape), # co
types.tensor(dtype, ret_shape),
) # h
