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 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 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 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 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 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 random_normal(RandomDistribution):
r"""
Returns a tensor with the specified shape, with random values from a normal
distribution.
Parameters
----------
shape: <K, i32> (Required)
* Target output tensor shape.
* ``K`` is the rank of the output tensor.
``shape[k] > 0`` for ``k = 0,..., K-1``.
mean: const<f32> (Optional)
The mean (center) of the normal distribution. Defaults to 0.0.
stddev: const<f32> (Optional)
The standard deviation (width) of the normal distribution. Defaults to ``1.0``.
seed: const<i32> (Optional)
Seed to create a reproducible sequence of values across multiple invokes.
Returns
-------
<\*, T>
* A tensor of the given target output shape filled with random values.
Attributes
----------
T: fp16, fp32
See Also
--------
random_categorical, random_bernoulli, random_uniform
"""
input_spec = (
InputSpec(
shape=IntTensorInputType(),
mean=FloatInputType(const=True, optional=True),
stddev=FloatInputType(const=True, optional=True),
seed=IntInputType(const=True, optional=True),
)
+ RandomDistribution.input_spec
)
def default_inputs(self):
return super().default_inputs() + \
DefaultInputs(
mean=0.,
stddev=1.,
seed=-1,
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
if self.mean.dtype != self.stddev.dtype:
raise ValueError("Incompatible primitive types in random_normal operation")
self.out_dtype = self.mean.dtype
return super().type_inference()
class random_bernoulli(RandomDistribution):
r"""
Returns a tensor with the specified shape, with random values from a Bernoulli
distribution.
.. math::
f(k) = \begin{cases}1-p &\text{if } k = 0\\
p &\text{if } k = 1\end{cases}
for :math:`k` in :math:`\{0, 1\}`.
Parameters
----------
shape: <K, i32> (Required)
* Target output tensor shape.
* ``K`` is the rank of the output tensor.
``shape[k] > 0`` for ``k = 0,..., K-1``.
prob: const<f32> (Optional)
* The probability of sampling ``1``. Defaults to ``0.5``.
seed: const<i32> (Optional)
* Seed to create a reproducible sequence of values across multiple invokes.
Returns
-------
<\*, T>
* A tensor of the given target output shape filled with random values.
Attributes
----------
T: fp16, fp32
See Also
--------
random_categorical, random_normal, random_uniform
"""
input_spec = (
InputSpec(
shape=IntTensorInputType(),
prob=FloatInputType(const=True, optional=True),
seed=IntInputType(const=True, optional=True),
)
+ RandomDistribution.input_spec
)
def default_inputs(self):
return super().default_inputs() + \
DefaultInputs(
seed=-1,
prob=0.5,
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
self.out_dtype = self.prob.dtype
return super().type_inference()
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 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 custom_topk(Operation):
input_spec = InputSpec(
x=TensorInputType(),
k=IntInputType(const=True, optional=True),
axis=IntInputType(const=True, optional=True),
sorted=BoolInputType(const=True, optional=True),
)
bindings = {
"class_name": "TopK",
"input_order": ["x"],
"parameters": ["k", "axis", "sorted"],
"description": "Top K Custom layer",
}
def default_inputs(self):
return DefaultInputs(
k=1,
axis=-1,
sorted=False,
)
def __init__(self, **kwargs):
super(custom_topk, self).__init__(**kwargs)
def type_inference(self):
x_type = self.x.dtype
x_shape = self.x.shape
k = self.k.val
axis = self.axis.val
if not is_symbolic(x_shape[axis]) and k > x_shape[axis]:
msg = "K={} is greater than size of the given axis={}"
raise ValueError(msg.format(k, axis))
ret_shape = list(x_shape)
ret_shape[axis] = k
return types.tensor(x_type, ret_shape), types.tensor(
types.int32, ret_shape)
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 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 softmax(Operation):
"""
Return ``exp(x) / tf.reduce_sum(tf.exp(x), axis)``.
Parameters
----------
x: tensor<\*?, T> (Required)
axis: const i32 (Optional)
* Default is ``-1``.
Returns
-------
tensor<\*?, fp32>
* A tensor of the same shape and type as ``x``.
Attributes
----------
T: fp32
"""
input_spec = InputSpec(
x=TensorInputType(),
axis=IntInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(axis=-1, )
def __init__(self, **kwargs):
super(softmax, self).__init__(**kwargs)
def type_inference(self):
return self.x.sym_type
@precondition(allow=VALUE)
def value_inference(self):
x = self.x.val
axis = self.axis.val
max_vals = np.max(x, axis=axis, keepdims=True)
temp = np.exp(x - max_vals)
return temp / np.sum(temp, axis=axis, keepdims=True)
class pixel_shuffle(Operation):
"""
Rearrange elements in a tensor from depth (channel) into spatial dimensions.
Equivalent to PyTorch's ``PixelShuffle``.
Parameters
----------
x: tensor<[n, C x f^2, H, W], T> (Required)
* Input tensor of rank ``4``.
upscale_factor: const<i32>
* Factor to increase spatial resolution by.
Returns
-------
tensor<[n, C, H x f, W x f], T>
* Where ``f`` is the upscale factor.
Attributes
----------
T: fp16, fp32
References
----------
`torch.nn.PixelShuffle <https://pytorch.org/docs/stable/generated/torch.nn.PixelShuffle.html?highlight=pixel%20shuffle#torch.nn.PixelShuffle>`_
"""
input_spec = InputSpec(
x=TensorInputType(),
upscale_factor=IntInputType(const=True),
)
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.upscale_factor.val
ret_shape = (n, c // (f * f), h * f, w * f)
return types.tensor(x_type, ret_shape)
class list_read(Operation):
"""
Read the value at location ``index`` of ``ls``.
Parameters
----------
ls: List[\*] (Required)
index: <i32> (Required)
* Size of the list.
Returns
-------
<\*,T>
* The element's value.
Attributes
----------
T: fp32, i32, bool
"""
input_spec = InputSpec(
ls=ListInputType(),
index=IntInputType(),
)
def __init__(self, **kwargs):
super(list_read, self).__init__(**kwargs)
def type_inference(self):
list_elem_type = self.ls.elem_type
if list_elem_type is None:
msg = ("Unknown element type. The List might not have been " +
"written to ({})")
raise ValueError(msg.format(self.name))
return list_elem_type
class topk(Operation):
"""
Returns a tensor containing top or bottom ``k`` values and the corresponding
indices of the input tensor along a given axis.
Parameters
----------
x: <\*?, T> (Required)
* Input tensor.
k: const<i32> (Optional)
* Default to ``1``.
* Number of values/indices to be computed along each axis.
* axis: const<i32> (Optional)
* Defaults to ``-1`` (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<\*?, T>
* Values of top/bottom ``k`` elements.
tensor<\*?, int32>
* Indices of the top/bottom ``k`` elements along axis.
Attributes
----------
T: fp32, int32
"""
input_spec = InputSpec(
x=TensorInputType(),
k=IntInputType(const=True, optional=True),
axis=IntInputType(const=True, optional=True),
ascending=BoolInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
k=1,
axis=-1,
ascending=False,
)
def __init__(self, **kwargs):
super(topk, self).__init__(**kwargs)
def type_inference(self):
x_type = self.x.dtype
x_shape = self.x.shape
k = self.k.val
axis = self.axis.val
if not is_symbolic(x_shape[axis]) and k > x_shape[axis]:
msg = "K={} is greater than size of the given axis={}"
raise ValueError(msg.format(k, axis))
ret_shape = list(x_shape)
ret_shape[axis] = k
return types.tensor(x_type, ret_shape), types.tensor(types.int32, ret_shape)
@precondition(allow=VALUE)
def value_inference(self):
indices = np.argsort(self.x.val, axis=self.axis.val)
if not self.ascending.val:
indices = np.argsort(-self.x.val, axis=self.axis.val)
slc = [slice(None)] * self.x.rank
slc[self.axis.val] = slice(0, self.k.val)
indices = indices[tuple(slc)]
values = np.take_along_axis(self.x.val, indices, axis=self.axis.val)
return values, indices
class scatter_along_axis(Operation):
"""
Scatter ``updates`` to ``data`` at locations ``indices`` along ``axis`` dimension
using ``mode`` operation.
Example: ``mode == update``.
* For ``i`` in ``[0, len(indices)]``:
.. math::
idx = indices[p_0, ..., p_{axis-1}, i, p_{axis+1}, ..., p_D]
.. math::
output[p_0, ..., p_{axis-1}, idx, p_{axis+1}, ..., p_D] =
.. math::
updates[p_0, ..., p_{axis-1}, i, p_{axis+1}, ..., p_D]
* For ``j! = i``:
.. math::
output[p_0, ..., p_{axis-1}, j, p_{axis+1}, ..., p_D] =
.. math::
data[p_0, ..., p_{axis-1}, j, p_{axis+1}, ..., p_D]
Example: ``mode == add``.
* For ``i`` in ``[0, len(indices)]``:
.. math::
idx = indices[p_0, ..., p_{axis-1}, i, p_{axis+1}, ..., p_D]
.. math::
output[p_0, ..., p_{axis-1}, idx, p_{axis+1}, ..., p_D] =
.. math::
updates[p_0, ..., p_{axis-1}, i, p_{axis+1}, ..., p_D] +
.. math::
x[p_0, ..., p_{axis-1}, indice[i], p_{axis+1}, ..., p_D]
* For ``j! = i``:
.. math::
output[p_0, ..., p_{axis-1}, j, p_{axis+1}, ..., p_D] =
.. math::
data[p_0, ..., p_{axis-1}, j, p_{axis+1}, ..., p_D]
Parameters
----------
data: tensor<\*D, T> (Required)
indices: tensor<\*K,T> (Required)
* ``rank(indices) == rank(data)``.
updates: tensor<\*K, T> (Required)
* Must be the same shape as ``indices``.
axis: const i32 (Optional)
* Default to ``0``.
mode: const string (Optional)
* Default to ``add``.
* Can be the following modes: ``update``, ``add``, ``sub``, ``mul``,
``div``, ``max``, ``min``.
Returns
-------
tensor<\*D, T>
* With the same type and shape as input ``x``.
Attributes
----------
U: fp16, fp32, i32
"""
input_spec = InputSpec(
data=TensorInputType(),
indices=IntTensorInputType(),
updates=TensorInputType(),
axis=IntInputType(const=True, optional=True),
mode=StringInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
axis=0,
mode="add",
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
@precondition(allow=VALUE)
def value_inference(self):
data = np.copy(self.data.val)
indices = self.indices.val
updates = self.updates.val
axis = self.axis.val
np_output = data
np.put_along_axis(np_output, indices, updates, axis=axis)
return np_output
def type_inference(self):
if self.axis.val < -self.data.rank or self.axis.val >= self.data.rank:
raise IndexError(
"Axis value {} is out of bounds for {} node {}".format(
self.axis.val, self.op_type, self.name))
axis = self.axis.val
axis = axis if axis >= 0 else axis + self.data.rank
assert is_compatible_symbolic_vector(self.indices.shape,
self.updates.shape)
assert self.data.rank == self.indices.rank
for i in range(self.data.rank):
if i != axis:
assert self.data.shape[i] == self.indices.shape[i]
return self.data.sym_type
class gather_along_axis(Operation):
"""
Take the values along ``axis`` at locations ``indices``.
.. math::
idx = indices[p_0, ..., p_{axis-1}, i, p_{axis+1}, ..., p_D]
.. math::
output[p_0, ..., p_{axis-1}, i, p_{axis+1}, ..., p_D] = = x[p_0, ..., p_{axis-1}, idx, p_{axis+1}, ..., p_D]
Parameters
----------
x: tensor<\*D, T> (Required)
indices: tensor<\*K, T> (Required)
* ``rank(indices) == rank(x)``.
axis: const i32 (Optional):
* Default to ``0``.
Returns
-------
tensor<\*D, T>:
* Output tensor has the same shape as ``indices``.
Attributes
----------
T: fp16, fp32, i32
"""
input_spec = InputSpec(
x=TensorInputType(),
indices=IntTensorInputType(),
axis=IntInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(axis=0, )
def __init__(self, **kwargs):
super().__init__(**kwargs)
@precondition(allow=VALUE)
def value_inference(self):
x = self.x.val
indices = self.indices.val
axis = self.axis.val
return np.take_along_axis(x, indices, axis)
def type_inference(self):
if self.x.rank != self.indices.rank:
raise ValueError("Rank mismatch between input and indices. \
Input rank: {}, indices rank: {}".format(
self.x.rank, self.indices.rank))
if self.axis.val < -self.x.rank or self.axis.val >= self.x.rank:
raise IndexError(
"Axis value {} is out of bounds for {} node {}".format(
self.axis.val, self.op_type, self.name))
axis = self.axis.val
axis = axis if axis >= 0 else axis + self.x.rank
for i in range(self.x.rank):
if i != axis:
assert self.x.shape[i] == self.indices.shape[i]
return types.tensor(self.x.dtype, self.indices.shape)
class gather(Operation):
"""
Gather slices from input ``x`` along dimension ``axis`` according to ``indices``,
similar to `tf.gather <https://www.tensorflow.org/api_docs/python/tf/gather>`_.
* If ``indices`` is scalar (0-D):
.. math::
output[p_0, ..., p_{axis-1}, ~~~~~~~~~~~~~~~~~~~~~~~~ p_{axis+1}, ..., p_{rank(x)-1}] =
.. math::
x[p_0, ..., p_{axis-1}, ~~~~~~~~~ indices, ~~~~~~~~ p_{axis+1}, ..., p_{rank(x)-1}]
Where ``rank(x)`` is the rank of ``x``. The ``output`` has rank ``rank(x) - 1``.
* If ``indices`` is 1-D tensor:
.. math::
output[p_0, ..., p_{axis-1}, ~~~~~~~~~~~~~ i, ~~~~~~~~~~~~~ p_{axis+1}, ..., p_{rank(*D)-1}] =
.. math::
x[p_0, ..., p_{axis-1}, ~~~~~~~~ indices[i], ~~~~~~~~ p_{axis+1}, ..., p_{rank(*D)-1}]
The output has rank ``rank(x)``.
* In general:
.. math::
output[p_0, ..., p_{axis-1}, ~~~~~~~~ i_0, ..., i_{M-1}, ~~~~~~~~ p_{axis+1}, ..., p_{rank(x)-1}] =
.. math::
x[p_0, ..., p_{axis-1}, ~~~~~~~ indices[i_0, ..., i_{M-1}], ~~~~~~~ p_{axis+1}, ..., p_{rank(x)-1}]
Where ``M = rank(indices)``.
Parameters
----------
x: tensor<\*D,T> (Required)
indices: tensor<\*N,i32> (Required)
* Indices values may be negative. More precisely, ``-D[axis]<= v < D[axis]`` for ``v`` in ``indices``.
axis: const i32 (Optional. Default=``0``)
* Negative axis is supported.
Returns
-------
tensor<\*K,T>
* Where ``K = D[:axis] + N + D[axis+1:]``.
Attributes
----------
U: fp16, fp32, i32
References
----------
See `tf.gather <https://www.tensorflow.org/api_docs/python/tf/gather>`_.
"""
input_spec = InputSpec(
x=TensorInputType(),
indices=IntInputType(),
axis=IntInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(axis=0, )
def __init__(self, **kwargs):
super().__init__(**kwargs)
@precondition(allow=VALUE | SYMBOL)
def value_inference(self):
x = self.x.sym_val
indices = self.indices.val
if indices is None:
# only allow x to be symbolic. indices cannot.
return None
scalar_indices = isinstance(indices, numbers.Integral)
axis = self.axis.val
if scalar_indices:
res = np.take(x, [indices], axis)
res2 = np.squeeze(res, axis=axis)
if isinstance(res2, np.ndarray) and len(res2.shape) == 0:
# res2 is a scalar, but represented as np.array(symbol,
# dtype=np.object) which np.squeeze can't remove.
return res2.item()
return res2
return np.take(x, indices, axis)
def type_inference(self):
out_type = self.x.dtype
if self.axis.val < -self.x.rank or self.axis.val >= self.x.rank:
raise IndexError(
"Axis value {} is out of bounds for {} node {}".format(
self.axis.val, self.op_type, self.name))
output_rank = self.x.rank - 1 + self.indices.rank
if output_rank == 0:
# output scalar
return out_type
axis = self.axis.val
axis = axis if axis >= 0 else axis + self.x.rank
out_shape = self.x.shape[:axis] + self.indices.shape + self.x.shape[
axis + 1:]
return types.tensor(out_type, out_shape)
class scatter(Operation):
"""
Scatter ``updates`` to ``data`` at locations ``indices`` at dimension ``axis``
by operation ``mode``.
Example: ``mode == update``.
* For ``i`` in ``[0, len(indices)]``:
.. math::
output[p_0, ..., p_{axis-1}, indice[i], p_{axis+1}, ..., p_D] =
.. math::
updates[p_0, ..., p_{axis-1}, i, p_{axis+1}, ..., p_D]
* For ``j != i``:
.. math::
output[p_0, ..., p_{axis-1}, j, p_{axis+1}, ..., p_D] =
.. math::
data[p_0, ..., p_{axis-1}, j, p_{axis+1}, ..., p_D]
Example: ``mode == add``.
* For ``i`` in ``[0, len(indices)]``:
.. math::
output[p_0, ..., p_{axis-1}, indice[i], p_{axis+1}, ..., p_D] =
.. math::
updates[p_0, ..., p_{axis-1}, i, p_{axis+1}, ..., p_D] +
.. math::
x[p_0, ..., p_{axis-1}, indice[i], p_{axis+1}, ..., p_D]
* For ``j != i``:
.. math::
output[p_0, ..., p_{axis-1}, j, p_{axis+1}, ..., p_D] =
.. math::
data[p_0, ..., p_{axis-1}, j, p_{axis+1}, ..., p_D]
Parameters
----------
data: tensor<\*D, T> (Required)
indices: tensor<[C],T> (Required)
* 1-D tensor.
updates: tensor<\*K, T> (Required)
* ``K = data.shape[:axis] + [len(indices)] + data.shape[axis+1:]``.
axis: const i32 (Optional)
* Default to ``0``.
mode: const string (Optional)
* Can be the following modes: ``update``, ``add``, ``sub``, ``mul``,
``div``, ``max``, ``min``.
* Default value is ``update``.
Returns
-------
tensor<\*D, T>
* With the same type and shape as input ``x``.
Attributes
----------
T: fp32
For example:
data = [[1, 2, 3], [4, 5, 6]]
indices = [1, 0]
updates = [[5, 6, 7], [8, 9, 10]]
axis = 0
mode = "update"
produces:
[[9, 11, 13], [9, 11, 13]]
"""
input_spec = InputSpec(
data=TensorInputType(),
indices=IntTensorInputType(),
updates=TensorInputType(),
axis=IntInputType(const=True, optional=True),
mode=StringInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
axis=0,
mode="add",
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
if self.axis.val < -self.data.rank or self.axis.val >= self.data.rank:
raise IndexError(
"Axis value {} is out of bounds for {} node {}".format(
self.axis.val, self.op_type, self.name))
axis = self.axis.val
axis = axis if axis >= 0 else axis + self.data.rank
expected_updates_shape = (self.data.shape[:axis] + self.indices.shape +
self.data.shape[axis + 1:])
err = "Updates shape {} is incorrect. It should be {}.".format(
self.updates.shape, expected_updates_shape)
assert is_compatible_symbolic_vector(
self.updates.shape, tuple(expected_updates_shape)), err
return self.data.sym_type
class batch_to_space(Operation):
"""
Rearrange elements in a tensor from batch into spatial dimension.
Parameters
----------
x: tensor<[n, C, H, W], T> (Required)
* Input tensor must have rank 4.
* The first and the second dimension are batch, channel, respectively
* The remaining dimensions (H, W) are treated as "spatial dimensions"
block_shape: const tensor<[2], i32> (Required)
* The length of the block_shape must be `2`
* It defines the shapes of the block in which the spatial dimensions are multiplied
crops: const tensor<[2, 2], i32> (Required)
* It must have shape `(2, 2)`
* It defines the amount to crop from each spatial dimensions
Returns
-------
tensor<[new_n, C, new_H, new_W], T>
* new_n = n / (block_shape[0] * block_shape[1])
* new_H = (H * block_shape[0]) - paddings[0][0] - padding[0][1]
* new_W = (W * block_shape[1]) - paddings[1][0] - padding[1][1]
* The output has the same rank as the input
Attributes
----------
T: fp16, fp32
"""
input_spec = InputSpec(
x=FloatTensorInputType(),
block_shape=IntInputType(const=True),
crops=IntInputType(const=True),
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
x_shape = self.x.shape
block_shape = self.block_shape.val
crops = self.crops.val
if self.x.rank != 4:
msg = "Input to batch_to_space op must be rank 4. Instead got an input with rank {}".format(
self.x.rank)
raise ValueError(msg)
if crops.shape != (block_shape.shape[0], 2):
msg = "block_shape and crops must have shape [2], [2, 2] accordingly in the batch_to_space op. "\
"Got {}, {}.".format(block_shape.shape, crops.shape)
raise ValueError(msg)
m = block_shape.shape[0]
if m != 2:
msg = "batch_to_space op only supports spatial dimensions = 2. Got {}".format(
m)
raise ValueError(msg)
b = x_shape[0]
c = x_shape[1]
spatial_shape = x_shape[2:2 + m]
if self.x.rank != m + 2:
raise ValueError("The input rank of batch_to_space op must exactly be " \
"len(block_shape){} + 2! Got {}".format(self.block_shape.val, self.x.rank))
if not is_symbolic(b) and b % np.prod(block_shape) != 0:
msg = (
"Batch size must be perfectly divided by the product of block_shape. Got batch size {}, and block_shape {}."
).format(b, block_shape)
raise ValueError(msg)
new_b = b / np.prod(block_shape)
new_spatial_shape = [
spatial_shape[i] * block_shape[i] for i in range(m)
]
cropped_spatial_shape = [
x - crops[i][0] - crops[i][1]
for i, x in enumerate(new_spatial_shape)
]
ret_shape = [new_b, c] + cropped_spatial_shape
x_type = self.x.dtype
return types.tensor(x_type, ret_shape)
class space_to_batch(Operation):
"""
Rearrange elements in a tensor from spatial into batch dimension.
Parameters
----------
x: tensor<[n, C, H, W], T> (Required)
* Input tensor must have rank 4.
* The first and the second dimension are batch, channel, respectively
* The remaining dimensions (H, W) are treated as "spatial dimensions"
block_shape: const tensor<[2], i32> (Required)
* The length of the block_shape must be `2`
* It defines the shapes of the block in which the spatial dimensions are divided
paddings: const tensor<[2, 2], i32> (Required)
* It must have shape `(2, 2)`
* It defines the padding for each spatial dimensions
Returns
-------
tensor<[new_n, C, new_H, new_W], T>
* new_n = n * block_shape[0] * block_shape[1]
* new_H = (H + paddings[0][0] + padding[0][1])/block_shape[0]
* new_W = (W + paddings[1][0] + padding[1][1])/block_shape[1]
* The output has the same rank as the input
Attributes
----------
T: fp16, fp32
"""
input_spec = InputSpec(
x=FloatTensorInputType(),
block_shape=IntInputType(const=True),
paddings=IntInputType(const=True),
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
x_shape = self.x.shape
block_shape = self.block_shape.val
paddings = self.paddings.val
if self.x.rank != 4:
msg = "Input to space_to_batch op must be rank 4. Instead got an input with rank {}".format(
self.x.rank)
raise ValueError(msg)
if paddings.shape != (block_shape.shape[0], 2):
msg = "block_shape and paddings must have shape [2], [2, 2] accordingly in the space_to_batch op. "\
"Got {}, {}.".format(block_shape.shape, paddings.shape)
raise ValueError(msg)
m = block_shape.shape[0]
if m != 2:
msg = "space_to_batch op only supports spatial dimensions = 2. Got {}".format(
m)
raise ValueError(msg)
b = x_shape[0]
c = x_shape[1]
spatial_shape = x_shape[2:2 + m]
if self.x.rank != m + 2:
raise ValueError("The input rank of space_to_batch op must exactly be " \
"len(block_shape){} + 2! Got {}".format(self.block_shape.val, self.x.rank))
padded_spatial_shape = [
x + paddings[i][0] + paddings[i][1]
for i, x in enumerate(spatial_shape)
]
new_b = b * np.prod(block_shape)
new_spatial_shape = [
padded_spatial_shape[i] / block_shape[i] for i in range(m)
]
ret_shape = [new_b, c] + new_spatial_shape
x_type = self.x.dtype
return types.tensor(x_type, ret_shape)
class cumsum(Operation):
"""
Returns the cumulative sum of the input along the given axis.
Parameters
----------
x: tensor<\*?, T> (Required)
* Input tensor.
axis: const<i32> (Optional)
* default to ``0``.
* Axis for which the cumulative sum is computed.
exclusive: const<bool> (Optional)
* Default to ``False``.
* When set to ``False``, inclusive cumsum is computed, that is the first element of
the output is identical to the first element in the input.
* When set to ``True``, exclusive cumsum is computed, which makes the first element
of output to ``0``.
reverse: const<bool> (Optional)
* Default to ``False``.
* When set to ``True``, perform cumsum in the reverse order.
Returns
-------
tensor<\*?, T>
* Same type and shape as the input tensor.
Attributes
----------
T: fp32, int32
"""
input_spec = InputSpec(
x=TensorInputType(),
axis=IntInputType(const=True, optional=True),
exclusive=BoolInputType(const=True, optional=True),
reverse=BoolInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
axis=0,
exclusive=False,
reverse=False)
def __init__(self, **kwargs):
super(cumsum, self).__init__(**kwargs)
@precondition(allow=VALUE)
def value_inference(self):
data = np.copy(self.x.val)
axis = self.axis.val
reverse = self.reverse.val
exclusive = self.exclusive.val
if reverse:
data = np.flip(data, axis=axis)
data = np.cumsum(data, axis=axis)
if exclusive:
zero_shape = np.copy(data.shape)
zero_shape[axis] = 1
data = np.concatenate((np.zeros(zero_shape, data)), axis=axis)
if reverse:
data = np.flip(data, axis=axis)
return data
def type_inference(self):
# Check range of axis
if self.axis.val < -1 or self.axis.val > self.x.rank - 1:
raise ValueError(
"axis should be in the range [-1, {}]".format(self.x.rank - 1)
)
return self.x.sym_type
class concat(Operation):
"""
Concatenates tensors along a dimension.
Parameters
----------
values: Tuple[tensor<[d0, d1, ..., d_axis_i, ..., d_n],T>] (Required)
* The number of dimensions of the input tensors must match, and all
dimensions except ``axis`` must be equal.
* The tensors may be variadic, but the number of tensors must be
determined at compile time (i.e. a tuple).
axis: const<int32> (Required)
* The dimension along which to concatenate. Must be in the range
``[-rank(values[i]), rank(values[i]))`` for all ``i``.
interleave: const<bool> (Optional, Default=False)
* If true, concatenate the inputs by interleaving them.
* If true, all the inputs to this op must have the exact same shape.
Examples
--------
.. sourcecode:: python
in1 : shape (3, 2), value = [[1, 2], [3, 4], [5, 6]]
in2 : shape (3, 2), value = [[7, 8], [9, 10], [11, 12]]
axis = 0
if interleave = False (default)
output : shape (6, 2)
output[0:3, :] = in1
output[3:6, :] = in2
value = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]]
if interleave = True
output : shape (6, 2)
output[0::2, :] = in1
output[1::2, :] = in2
value = [[1, 2], [7, 8], [3, 4], [9, 10], [5, 6], [11, 12]]
Returns
-------
tensor<[d0, d1,...d_axis_out, ..., d_n],T>
* Where ``d_axis_out = sum(d_axis_i)``.
Attributes
----------
T: fp32, int32
"""
input_spec = InputSpec(values=TupleInputType(),
axis=IntInputType(const=True),
interleave=BoolInputType(const=True,
optional=True))
def default_inputs(self):
return DefaultInputs(
interleave=False,
)
def __init__(self, **kwargs):
super(concat, self).__init__(**kwargs)
def type_inference(self):
concat_dim_len = 0
if len(self.values) == 0:
raise ValueError("Concat {} got 0 values".format(self.name))
# Validate values have the same rank
rank = self.values[0].rank
for v in self.values:
if v.rank != rank:
msg = "Input {} has rank {} != other inputs rank {}"
raise ValueError(msg.format(v.name, v.rank, rank))
# Check concat axis is within (-rank, rank)
concat_axis = self.axis.val
if concat_axis < 0:
concat_axis += rank
if rank > 0 and (concat_axis < 0 or concat_axis >= rank):
msg = "In {} of op_type {}: axis out of bound for input " + "(rank {})"
raise ValueError(msg.format(self.name, self.op_type, rank))
# Validate primitive types are compatible
dtype = self.values[0].dtype
for v in self.values[1:]:
new_dtype = promoted_primitive_type(v.dtype, dtype)
if new_dtype is None:
msg = "Incompatible primitive types concat: {} vs {}"
raise ValueError(msg.format(v.dtype, dtype))
dtype = new_dtype
# validate that non-axis dimensions match
retshape = list(self.values[0].shape)
for v in self.values[1:]:
for i in range(rank):
if is_symbolic(retshape[i]) or is_symbolic(v.shape[i]):
continue
if i != concat_axis and retshape[i] != v.shape[i]:
msg = 'Dimension mismatch in {} ("{}"): shapes {} vs. {}'
raise ValueError(
msg.format(self.op_type, self.name, retshape, v.shape)
)
if self.interleave.val and retshape[i] != v.shape[i]:
msg = 'Dimension mismatch in {} ("{}"): shapes {} vs. {}. ' \
'All inputs must have same shape when \'interleave\' option is True.'
raise ValueError(
msg.format(self.op_type, self.name, retshape, v.shape)
)
# Get length of concat dim
concat_dim_len = 0
for v in self.values:
if len(v.shape) == 0:
taxis = 1
else:
taxis = v.shape[concat_axis]
if is_symbolic(taxis):
concat_dim_len = get_new_symbol()
break
concat_dim_len += taxis
if len(retshape) == 0:
retshape = [concat_dim_len]
else:
retshape[concat_axis] = concat_dim_len
return types.tensor(dtype, retshape)
@precondition(allow=VALUE | SYMBOL | NONE)
def value_inference(self):
values = []
for v in self.values:
if v.sym_val is not None:
values.append(v.sym_val)
continue
if v.rank == 0:
values.append(get_new_symbol())
continue
if any_symbolic(v.shape):
values.append(None)
continue
# we support value inference when number of elements for each tensor is less than 10
shape = v.shape
num_element = np.prod(shape)
if num_element > 10:
values.append(None)
continue
symbolic_tensor = [get_new_symbol() for _ in range(num_element)]
symbolic_tensor = np.reshape(np.array(symbolic_tensor), shape)
values.append(symbolic_tensor)
if any([val is None for val in values]):
return None
if not isinstance(values[0], np.ndarray) or values[0].shape == ():
return np.stack(values, axis=self.axis.val)
return np.concatenate(values, axis=self.axis.val)
class stack(Operation):
"""
Concatenates tensors along a dimension.
Parameters
----------
values: Tuple[tensor<[d0, d1,...d_axis_i, ..., d_n],T>] (Required)
* All tensors must have identical shape.
axis: const<i32> (Required)
* The dimension along which to concatenate. Must be in the range ``[-rank(values[i]), rank(values[i]))`` for all ``i``.
Returns
-------
tenor<[d0, d1,...d_axis_out, ..., d_n],T>
* Where ``d_axis_out = sum(d_axis_i)``.
Attributes
----------
T: fp32
"""
input_spec = InputSpec(values=TupleInputType(),
axis=IntInputType(const=True),)
def __init__(self, **kwargs):
super(stack, self).__init__(**kwargs)
def type_inference(self):
num_tensors = len(self.values)
if num_tensors == 0:
raise ValueError("Cannot stack 0 tensor")
# get the first value without symbolic shape
t_shape = None
for value in self.values:
if not any_symbolic(value.shape):
t_shape = value.shape
break
t_shape = self.values[0].shape if t_shape is None else t_shape
# compare all shape
for t in self.values:
if not is_compatible_symbolic_vector(t.shape, t_shape):
msg = "Component tensor {} has shape {}, others have {}"
raise ValueError(msg.format(t.name, t.shape, t_shape))
ret_shape = list(t_shape)
ret_shape.insert(self.axis.val, num_tensors)
return types.tensor(self.values[0].dtype, ret_shape)
@precondition(allow=VALUE | SYMBOL | NONE)
def value_inference(self):
is_all_rank_zero = all([v.rank == 0 for v in self.values])
values = [
v.sym_val if v.sym_val is not None else get_new_symbol()
for v in self.values
]
if any([is_symbolic(v) for v in values]) and not is_all_rank_zero:
return None
return np.stack(values, self.axis.val)
class random_uniform(RandomDistribution):
r"""
Returns a tensor with the specified shape with random values from a uniform
distribution. Samples are uniformly distributed over the half-open interval
``[low, high)`` (includes low, but excludes high).
.. math::
p(x) = \frac{1}{high - low}
For a real number :math:`x`.
When ``high == low``, values of ``low`` will be returned. If ``high < low``,
the results are officially undefined and may eventually raise an error.
Parameters
----------
shape: <K, i32> (Required)
* Target output tensor shape.
* ``K`` is the rank of the output tensor.
``shape[k] > 0`` for ``k = 0,..., K-1``.
low: const<f32> (Optional)
* Lower boundary of the output interval (inclusive). Defaults to ``0.0``.
high: const<f32> (Optional)
* Upper boundary of the output interval (exclusive). Defaults to ``1.0``.
seed: const<i32> (Optional)
* Seed to create a reproducible sequence of values across multiple invokes.
Returns
-------
<\*, T>
* A tensor of the given target output shape filled with random values.
Attributes
----------
T: fp16, fp32
See Also
--------
random_categorical, random_bernoulli, random_normal
"""
input_spec = (
InputSpec(
shape=IntTensorInputType(),
low=FloatInputType(const=True, optional=True),
high=FloatInputType(const=True, optional=True),
seed=IntInputType(const=True, optional=True),
)
+ RandomDistribution.input_spec
)
def default_inputs(self):
return super().default_inputs() + \
DefaultInputs(
low=0.,
high=1.,
seed=-1,
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
if self.low.dtype != self.high.dtype:
raise ValueError("Incompatible primitive types in random_uniform operation")
self.out_dtype = self.low.dtype
return super().type_inference()
class non_maximum_suppression(Operation):
"""
Applies non-maximum suppression (NMS) on the input box coordinates according
to their intersection-over-union (IoU).
NMS selects a subset of bounding boxes in descending order of score, and removes
boxes that have high intersection-over-union (IOU) overlap with previously-selected
boxes.
Parameters
----------
boxes: tensor<[n, B, 4], T> (Required)
* Box coordinates on which to perform NMS.
scores: tensor<[n, B, K], T> (Required)
* Scores for each one of the boxes.
iou_threshold: const<T> (Required)
* The intersection over union (``IoU``) threshold over which boxes are
suppressed. NMS remove all overlapping boxes with ``IoU > iou_threshold``.
score_threshold: const<T> (Required)
* Before IoU suppression is performed, boxes with class scores below this
threshold are rejected.
max_boxes: const<i32> (Required)
* Maximum number of boxes to select. If the number of surviving boxes are
less, output is padded up to this number.
per_class_suppression: const<bool> (Optional)
* Default to ``False``.
* If ``True``, suppression is performed independently within boxes of each class.
Returns
-------
tensor<[n, max_boxes, 4], T>
* Coordinates of selected boxes.
tensor<[n, max_boxes, K], T>
* Scores of selected boxes.
tensor<[n, max_boxes], i32>
* Indices of selected boxes.
tensor<[n], i32>
* Number of boxes selected for each batch.
Attributes
----------
T: fp32
"""
input_spec = InputSpec(
boxes=TensorInputType(),
scores=TensorInputType(),
iou_threshold=FloatInputType(const=True),
score_threshold=FloatInputType(const=True),
max_boxes=IntInputType(const=True),
per_class_suppression=BoolInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
per_class_suppression=False)
def __init__(self, **kwargs):
super(non_maximum_suppression, self).__init__(**kwargs)
def type_inference(self):
boxes_dtype = self.boxes.dtype
scores_dtype = self.scores.dtype
n_batch, _, n_score = self.scores.shape
max_boxes = self.max_boxes.val
return (
types.tensor(boxes_dtype, (n_batch, max_boxes, 4)),
types.tensor(scores_dtype, (n_batch, max_boxes, n_score)),
types.tensor(types.int32, (n_batch, max_boxes)),
types.tensor(types.int32, (n_batch,)),
)
class one_hot(Operation):
"""
Returns one-hot vectors whose locations represented in ``indices`` take the ``on_value``,
while other locations take the ``off_value``.
Parameters
----------
indices: tensor<[D],T> (Required)
* Tensor, values indicate the locations for each one-hot vector to take the ``on_value``.
one_got_vector_size: i32 (Required)
* Indicates the number of returning vectors.
axis: const i32 (Optional)
* Indicates which dimension to append the new axis.
* If the input indices is rank ``D``, the output tensor will have rank ``D+1``.
* Default to ``-1`` (the last dimension).
on_value: const i32 (Optional)
* Values for locations where defined in ``indices``.
* Default to ``1``.
off_value: const i32 (Optional)
* Default to ``0``.
Returns
-------
tensor<\*?,T>
* A tensor that contains one-hot vectors.
Attributes
----------
T: fp32
"""
input_spec = InputSpec(
indices=IntTensorInputType(),
one_hot_vector_size=IntInputType(),
axis=IntInputType(const=True, optional=True),
on_value=IntOrFloatInputType(const=True, optional=True),
off_value=IntOrFloatInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
axis=-1,
on_value=1,
off_value=0,
)
def __init__(self, **kwargs):
super(one_hot, self).__init__(**kwargs)
def type_inference(self):
on_type = self.on_value.dtype
off_type = self.off_value.dtype
if on_type != off_type:
raise TypeError(
"Parameters on_value and off_value must have same input types."
)
if self.axis.val < -self.indices.rank - 1 or self.axis.val > self.indices.rank:
raise IndexError(
"Axis value {} is out of bounds for {} node {}".format(
self.axis.val, self.op_type, self.name
)
)
indices_shape = list(self.indices.shape)
depth_value = self.one_hot_vector_size.val
if depth_value is None:
depth_value = get_new_symbol()
elif depth_value < 0:
raise ValueError("Parameter one_hot_vector_size must be non-negative")
retshape = indices_shape
if self.axis.val < 0:
cut = len(retshape) + self.axis.val + 1
else:
cut = self.axis.val
retshape = retshape[0:cut] + [depth_value] + retshape[cut:]
return types.tensor(on_type, retshape)
