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)