class scatter_nd(Operation):
"""
Scatter ``updates`` to ``data`` at locations ``indices``.
The ``indices`` is a K-dim tensor, where ``indices[i_0,...,i_{K-2}]`` defines a
slice of ``data``, ``K = rank(indices)``, and ``data[indices[i_0, ..., i_{K-2}]]``
has rank ``rank(data) - indices.shape[-1]``.
* Example: ``mode == update``: The ``output`` is set to ``data`` initially, and
the op updates ``output`` as follows:
.. math::
output[indices[i_0, ..., i_{K-2}]]= updates[indices[i_0, ..., i_{K-2}]]
* Example: ``mode == add``. The update rule is:
.. math::
output[indices[i_0, ..., i_{K-2}]] += updates[indices[i_0, ..., i_{K-2}]]
Parameters
----------
data: tensor<\*D,T> (Required)
indices: tensor<\*K,i32> (Required)
updates: tensor<\*K, T> (Required)
* Must be the shape as ``K[:-1]+data.shape[K[-1]:]``.
mode: const string (Optional)
* Default to ``add``.
* Can be the following modes: ``update``, ``add``, ``sub``, ``mul``,
``div``, ``max``, ``min``.
Returns
-------
tensor<\*D,T>
* A tensor with the same shape and type as ``data``.
Attributes
----------
T: fp16, fp32, i32
"""
input_spec = InputSpec(
data=TensorInputType(),
indices=IntTensorInputType(),
updates=TensorInputType(),
mode=StringInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(mode="add", )
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
assert self.indices.shape[-1] <= self.data.rank
expected_updates_shape = (self.indices.shape[:-1] +
self.data.shape[self.indices.shape[-1]:])
assert is_compatible_symbolic_vector(self.updates.shape,
tuple(expected_updates_shape))
return self.data.sym_type
class 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 Pooling(Operation):
"""
Pooling Op Superclass
"""
input_spec = InputSpec(
x=TensorInputType(),
kernel_sizes=IntTensorInputType(const=True),
strides=IntTensorInputType(const=True, optional=True),
pad_type=StringInputType(const=True),
pad=IntTensorInputType(const=True, optional=True),
ceil_mode=BoolInputType(const=True, optional=True),
)
def default_inputs(self):
num_spatial_dims = self.x.rank - 2
return DefaultInputs(
strides=[1]*num_spatial_dims,
pad=[0]*2*num_spatial_dims,
ceil_mode=False,
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
ksize = self.kernel_sizes.val
x_shape = self.x.shape
D_in_rank = len(x_shape) - 2
strides = [1] * D_in_rank if self.strides is None else self.strides.val
pad_type = "valid" if self.pad_type is None else self.pad_type.val.lower()
if pad_type not in ["valid", "same", "custom"]:
raise ValueError("Unrecognized value of pad_type : {}".format(pad_type))
pad = None if self.pad is None else self.pad.val
D_in = x_shape[2:] # spatial dimensions
if self.ceil_mode.val:
if D_in_rank > 2:
raise ValueError('pool: ceil_mode only supported for 1D or 2D pool')
if pad_type == "same" and self.ceil_mode.val:
raise ValueError("ceil_mode must be False when pad_type==same")
if pad is not None:
for i in range(D_in_rank):
if pad[2*i] != pad[2*i+1]:
raise ValueError("Padding must be symmetric if ceil_mode is True")
D_out_shape = spatial_dimensions_out_shape(
pad_type=pad_type,
input_shape=D_in,
kernel_shape=ksize,
strides=strides,
custom_pad=pad,
ceil_mode=self.ceil_mode.val,
)
ret_shape = list(x_shape[:2]) + D_out_shape
return types.tensor(self.x.dtype, tuple(ret_shape))
class resample(_resample_iOS15):
"""
iOS16 version of resample supports float16 coordinates
"""
input_spec = InputSpec(
x=TensorInputType(),
coordinates=ScalarOrTensorInputType(type_domain=(np.int32, np.float32, np.float16)),
sampling_mode=StringInputType(const=True),
padding_mode=StringInputType(const=True),
padding_value=FloatInputType(const=True),
coordinates_mode=StringInputType(const=True),
align_corners=BoolInputType(const=True),
)
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
return super().type_inference()
class 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 cast(Operation):
"""
Cast the input ``x`` to the new type ``dtype``.
Parameters
----------
x: tensor<[\*d], T> (Required)
dtype: const str (Required)
* Can be one of the following types: ``int32``, ``int64``, ``fp32``, ``fp64``.
Returns
-------
tensor<[\*d], dtype>
* A tensor of the same shape as ``x``, with type ``dtype``.
Attributes
----------
T: i32, i64, fp16, fp32, fp64, bool.
"""
input_spec = InputSpec(x=ScalarOrTensorInputType(),
dtype=StringInputType(const=True))
def __init__(self, **kwargs):
super().__init__(**kwargs)
def type_inference(self):
type_map = {
"int32": types.int32,
"int64": types.int64,
"fp16": types.fp16,
"fp32": types.fp32,
"fp64": types.fp64,
"bool": types.bool,
}
if self.dtype.val not in type_map.keys():
raise NotImplementedError(
"Parameter dtype of the cast operation can be one of the {}. "
"Provided {}".format(type_map.keys(), self.dtype.val))
if not types.is_tensor(self.x.sym_type):
return type_map[self.dtype.val]
ret_shape = self.x.shape
return types.tensor(type_map[self.dtype.val], ret_shape)
@precondition(allow=VALUE | SYMBOL)
def value_inference(self):
return self.get_cast_value(self.x, self.dtype.val)
@staticmethod
def get_cast_value(input_var, dtype_val):
type_map = {
"int32": np.int32,
"int64": np.int64,
"fp16": np.float16,
"fp32": np.float32,
"fp64": np.float64,
"bool": np.bool,
}
if dtype_val not in type_map.keys():
raise NotImplementedError(
"Parameter dtype of the cast operation can be one of the {}. "
"Provided {}".format(type_map.keys(), dtype_val))
if input_var.val is None:
if input_var.sym_val is not None and not is_symbolic(
input_var.sym_val) and len(input_var.sym_val.shape) == 1:
result = [
np.array(val).astype(dtype=type_map[dtype_val]).item()
if not is_symbolic(val) else val
for val in input_var.sym_val
]
return np.array(result)
return None
if not types.is_tensor(input_var.sym_type):
return input_var.val.astype(dtype=type_map[dtype_val])
else:
return np.array(input_var.val).astype(dtype=type_map[dtype_val])
class gelu(Operation):
"""
Return the elementwise Gaussian error linear unit activation function for ``x``.
You can use ``EXACT``, ``TANH_APPROXIMATION``, or ``SIGMOID_APPROXIMATION`` values
based on the following formulas:
* ``EXACT``:
.. math::
f(x) = 0.5x\\left ( 1+\\rm{erf}\\left ( \\frac{x}{\\sqrt{2}} \\right ) \\right )
* ``TANH_APPROXIMATION``:
.. math::
f(x) = 0.5x\\left ( 1+\\rm{tanh}\\left ( \\sqrt{2/\\pi}\\left ( x + 0.044715x^3 \\right ) \\right ) \\right )
* ``SIGMOID_APPROXIMATION``:
.. math::
f(x) = x*\\rm{sigmoid}(1.702x)
Parameters
----------
x: tensor<\*?, T> (Required)
mode: const str (Optional)
* Use ``'EXACT'``, ``'TANH_APPROXIMATION'``, or ``'SIGMOID_APPROXIMATION'`` for ``str``.
* Default is ``'EXACT'``.
Returns
-------
tensor<\*?, T>
* A tensor of the same shape and type as ``x``.
Attributes
----------
T: fp32
"""
input_spec = InputSpec(
x=ScalarOrTensorInputType(),
mode=StringInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(mode="EXACT", )
def __init__(self, **kwargs):
super(gelu, self).__init__(**kwargs)
@precondition(allow=VALUE)
def value_inference(self):
if self.mode.val == "TANH_APPROXIMATION":
a = np.sqrt(
2 / np.pi) * (self.x.val + 0.044715 * np.power(self.x.val, 3))
return 0.5 * self.x.val * (1 + np.tanh(a))
elif self.mode.val == "SIGMOID_APPROXIMATION":
return self.x.val * (1 / (1 + np.exp(-(1.702 * self.x.val))))
else:
sqaure_root_of_2 = np.sqrt(2)
vfunc = np.vectorize(lambda x: 0.5 * x *
(1 + math.erf(x / sqaure_root_of_2)))
return vfunc(self.x.val)
def type_inference(self):
allowed_values = {
"EXACT", "TANH_APPROXIMATION", "SIGMOID_APPROXIMATION"
}
if self.mode.val not in allowed_values:
msg = '"gelu" op: unrecognized value of mode: "{}". Allowed values are {}'
raise ValueError(msg.format(self.mode.val, allowed_values))
return self.x.sym_type
class cast(Operation):
"""
Cast the input ``x`` to the new type ``dtype``.
Parameters
----------
x: tensor<[\*d], T> (Required)
dtype: const str (Required)
* Can be one of the following types: ``int32``, ``int64``, ``fp32``, ``fp64``.
Returns
-------
tensor<[\*d], dtype>
* A tensor of the same shape as ``x``, with type ``dtype``.
Attributes
----------
T: i32, i64, fp32, fp64, bool.
"""
input_spec = InputSpec(x=ScalarOrTensorInputType(),
dtype=StringInputType(const=True))
def __init__(self, **kwargs):
super(cast, self).__init__(**kwargs)
def type_inference(self):
type_map = {
"int32": types.int32,
"int64": types.int64,
"fp32": types.fp32,
"fp64": types.fp64,
"bool": types.bool,
}
if self.dtype.val not in type_map.keys():
raise NotImplementedError(
"Parameter dtype of the cast operation can be one of the {}. "
"Provided {}".format(type_map.keys(), self.dtype.val))
if not types.is_tensor(self.x.sym_type):
return type_map[self.dtype.val]
ret_shape = self.x.shape
return types.tensor(type_map[self.dtype.val], ret_shape)
@precondition(allow=VALUE)
def value_inference(self):
type_map = {
"int32": np.int32,
"int64": np.int64,
"fp32": np.float32,
"fp64": np.float64,
"bool": np.bool,
}
if self.dtype.val not in type_map.keys():
raise NotImplementedError(
"Parameter dtype of the cast operation can be one of the {}. "
"Provided {}".format(type_map.keys(), self.dtype.val))
if not types.is_tensor(self.x.sym_type):
return self.x.val.astype(dtype=type_map[self.dtype.val])
else:
return np.array(self.x.val).astype(dtype=type_map[self.dtype.val])
class rnn(Operation):
"""
Recurrent neural network (RNN).
.. math::
h_t = activation(W_{ih} x_t + b_{ih} + W_{hh} h_(t−1) + b_{hh})
Where:
* ``W_{ih}`` is input weight.
* ``W_{hh}`` is hidden/recurrent weight.
* ``h_t`` is the hidden state at time ``t``.
* ``x_t`` is the input at time ``t``.
* ``h_(t-1)`` is the hidden state of the layer at time ``t-1`` or the initial
hidden state at time ``0``.
Parameters
----------
x: <s, b, I, T> (Required)
* ``s`` is the sequence length, ``b`` is the batch size, and ``I`` is the
input dimension.
initial_h: <b, H, T> (Required)
* ``H`` denotes hidden size.
weight_ih: const<H, I, T> (Required) - Input-hidden weight matrix
weight_hh: const<H, H, T> (Required) - Hidden-hidden weight matrix
bias: const<H, T> (Optional) [Default all 0s]
* bias for input-hidden and hidden-hidden
direction: const<str> (Optional) [Default=forward]
* Either ``forward`` or ``reverse``.
output_sequence: const<bool> (Optional) [Default=False]
* Outputs every step if ``True``.
activation: const<str> (Optional) [Default=tanh]
* Supported activation functions: ``relu``, ``tanh``, ``sigmoid``,
``sigmoid_hard``, ``scaled_tanh``, and ``linear``.
Returns
-------
<s, b, H, T> or <1, b, H, T>
* If ``output_sequence == True`` (hidden states from every step):
``<s, b, H, T>``.
* Else ``<1, b, H, T>`` (hidden states of the final step).
<b, H, T>
* Hidden states of the final step.
Attributes
----------
T: fp32
"""
input_spec = InputSpec(
x=TensorInputType(),
initial_h=TensorInputType(),
weight_ih=TensorInputType(const=True),
weight_hh=TensorInputType(const=True),
bias=TensorInputType(const=True, optional=True),
direction=StringInputType(const=True, optional=True),
output_sequence=BoolInputType(const=True, optional=True),
activation=StringInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(bias=None,
direction="forward",
output_sequence=False,
activation="tanh")
def __init__(self, **kwargs):
super(rnn, self).__init__(**kwargs)
def type_inference(self):
if self.x.rank != 3:
raise ValueError(
"Invalid input shape. Expecting Rank 3 input, got {}".format(
len(self.x.rank)))
sequence_length, batch_size, input_size = self.x.shape
if self.weight_ih.rank != 2 or self.weight_hh.rank != 2:
raise ValueError(
"Invalid weight shape. Expecting Rank 2 input, got weight_ih {}, weight_hh {}"
).format(self.weight_ih.rank, self.weight_hh.rank)
hidden_size, _ = self.weight_ih.shape
direction = self.direction.val
valid_directions = {"forward", "reverse"}
if direction not in valid_directions:
raise ValueError(
"Direction {} not supported. Supported directions: {}".format(
direction, valid_directions))
out_seq_len = sequence_length if self.output_sequence.val else 1
output_shape = [out_seq_len, batch_size, hidden_size]
output_h_shape = [batch_size, hidden_size]
return (
types.tensor(self.x.dtype, tuple(output_shape)),
types.tensor(self.x.dtype, tuple(output_h_shape)),
)
class lstm(Operation):
r"""
Single long short-term memory (LSTM) sequence.
.. math::
i_t = \rm{recurrent\_activation}(W_{ii} x_t + B_{ii} + W_{hi} h_(t-1) + B_{hi})
.. math::
f_t = \rm{recurrent\_activation}(W_{if} x_t + B_{if} + W_{hf} h_(t-1) + B_{hf})
.. math::
z_t = cell_activation(W_{iz} x_t + B_{iz} + W_{hz} h_(t-1) + B_{hz})
.. math::
o_t = \rm{recurrent\_activation}(W_{io} x_t + B_{io} + W_{ho} h_(t-1) + B_{ho})
.. math::
c_t = f_t * c_(t-1) + i_t * z_t
.. math::
h_t = o_t * activation(c_t)
Where:
* ``i_t``, ``f_t``, ``o_t``, and ``z_t`` are input, forget, output, and cell gates,
respectively, at time ``t``.
* ``c_t`` is cell state at time ``t``.
* ``h_t`` is the hidden state at time ``t``.
* ``W_{ii}``, ``W_{if}``, ``W_{io}``, and ``W_{iz}`` are input weights for input,
forget, output and cell gate, respectively.
* ``W_{hi}``, ``W_{hf}``, ``W_{ho}``, and ``W_{hz}`` are recurrent weights for input,
forget, output and cell gate, respectively.
Parameters
----------
x: <s, b, I, T> (Required)
* ``s`` is the sequence length, ``b`` is the batch size, and ``I`` is the
input dimension.
initial_h: <b, DIRECTION*H, T> (Required)
* Initial hidden state. ``DIRECTION = 1`` for uni-directional, ``2`` for
bi-directional LSTM.
* ``H`` denotes hidden size.
* ``[b, :H]`` and ``[b, H:]`` represents forward and reverse direction
values, respectively.
initial_c: <b, DIRECTION*H, T> (Required)
* Initial cell state.
* Format is same as ``initial_h``.
weight_ih: const<4*H, I, T> (Required)
* Input-hidden weight matrix
* Weight tensor should be in order of
``[input_gate, forget_gate, output_gate, cell_gate]``.
* If direction=="bidirectional", this is applied in forward direction.
* If direction=="forward" or "backward" these weights are used.
weight_hh: const<4*H, H, T> (Required)
* Hidden-hidden weight matrix.
* Weight tensor should be in order of
``[input_gate, forget_gate, output_gate, cell_gate]``.
* If direction=="bidirectional", this is applied in forward direction.
* If direction=="forward" or "backward" these weights are used.
bias: const<4*H, T> (Optional) [Default all 0s]
* bias = input-hidden bias + hidden-hidden bias
* If direction=="bidirectional", this is applied in forward direction.
* If direction=="forward" or "backward" this bias are used.
peephole: const<3*H, T> (Optional, default to 0)
* Weight tensor for peephole.
* Order is ``[input_gate, forget_gate, output_gate]``.
* Shape of each peephole vector is ``(H,)`` (``H`` is hidden size).
* If direction=="bidirectional", this is applied in forward direction.
* If direction=="forward" or "backward" these weights are used.
weight_ih_back: const<4*H, I, T> (Optional) -
* Input-hidden weight matrix for backward direction for `bidirectinal LSTM`.
* Weight tensor should be in order of
``[input_gate, forget_gate, output_gate, cell_gate]``.
* Must be provided for `bidirectional LSTM`.
* This is only used when `direction` is "bidirectional".
* For direction="reverse" use `weight_ih` instead.
weight_hh_back: const<4*H, H, T> (Optional) - Hidden-hidden weight matrix
* Hidden-hidden weight matrix for backward direction for `bidirectinal LSTM`.
* Weight tensor should be in order of
``[input_gate, forget_gate, output_gate, cell_gate]``.
* Must be provided for `bidirectional LSTM`.
* This is only used when `direction` is "bidirectional".
* For direction="reverse" use `weight_hh` instead.
bias_back: const<4*H, T> (Optional) [Default all 0s]
* bias = input-hidden bias + hidden-hidden bias.
* Bias of backward direction for `bidirectional lstm`
* This is only used when `direction` is "bidirectional".
* For direction="reverse" use `bias` instead.
peephole_back: const<3*H, T> (Optional, default to 0)
* Weight tensor for peephole in backward direction for `bidirectional LSTM`.
* Order is ``[input_gate, forget_gate, output_gate]``.
* Shape of each peephole vector is ``(H,)`` (``H`` is hidden size).
* Peephole of backward direction for `bidirectional lstm`
* Bias of backward direction for `bidirectional lstm`
* This is only used when `direction` is "bidirectional".
* For direction="reverse" use `peephole` instead.
direction: const<str> (Optional) [Default=forward]
* One of the following: ``forward``, ``reverse``, or ``bidirectional``.
* Must match ``DIRECTIONAL`` in initial states and weight parameters.
output_sequence: const<bool> (Optional) [Default=False]
* Outputs every step if ``True``.
recurrent_activation: const<str> (Optional) [Default=sigmoid]
* Activation applied on input, forget, and output gates.
cell_activation: const<str> (Optional) [Default=tang]
* Activation applied on cell gate.
activation: const<str> (Optional) [Default=tanh]
* Activation applied on output gate.
clip: const<fp32> (optional) [Default=None]
* Cell gate is clipped to ``[-clip, +clip]``.
Returns
-------
<s, b, DIRECTION*H, T> or <1, b, DIRECTION*H, T>
* If ``output_sequence == True`` (hidden states from every step):
``<s, b, DIRECTION*H, T>``.
* Else ``<1, b, DIRECTION*H, T>`` (hidden states of the final step).
<b, DIRECTION*H, T>
* Hidden states of the final step.
<b, DIRECTION*H, T>
* Memory state of the final step.
Attributes
----------
T: fp32
"""
input_spec = InputSpec(
x=TensorInputType(),
initial_h=TensorInputType(),
initial_c=TensorInputType(),
weight_ih=TensorInputType(const=True), # ifoz layout,
weight_hh=TensorInputType(const=True), # ifoz layout
bias=TensorInputType(const=True, optional=True), # ifoz layout
peephole=TensorInputType(const=True, optional=True), # ifo layout
weight_ih_back=TensorInputType(const=True,
optional=True), # ifoz layout,
weight_hh_back=TensorInputType(const=True,
optional=True), # ifoz layout
bias_back=TensorInputType(const=True, optional=True), # ifoz layout
peephole_back=TensorInputType(const=True, optional=True), # ifo layout
direction=StringInputType(const=True, optional=True),
output_sequence=BoolInputType(const=True, optional=True),
recurrent_activation=StringInputType(const=True, optional=True),
cell_activation=StringInputType(const=True, optional=True),
activation=StringInputType(const=True, optional=True),
clip=FloatInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(bias=None,
direction="forward",
output_sequence=False,
recurrent_activation="sigmoid",
cell_activation="tanh",
activation="tanh",
peephole=None,
clip=None)
def __init__(self, **kwargs):
super(lstm, self).__init__(**kwargs)
def type_inference(self):
if self.x.rank != 3:
raise ValueError(
"Invalid input shape. Expecting Rank 3 input, got {}".format(
len(self.x.rank)))
sequence_length, batch_size, input_size = self.x.shape
def weight_shape_check(wt_ih, wt_hh):
if wt_ih.rank != 2 or wt_hh.rank != 2:
raise ValueError(
"Expecting Rank 2 input, got weight_ih rank: {}, weight_hh rank: {}"
).format(wt_ih.rank, wt_hh.rank)
hidden_size = wt_hh.shape[1]
if wt_hh.shape[0] // hidden_size != 4 or wt_ih.shape[
0] // hidden_size != 4:
raise ValueError(
"Incorrect weight matrix: hidden dim size mismatch. \
Provided weight_ih {}, weight_hh {}. Expecting <4*H, H>"
).format(wt_ih.shape, wt_hh.shape)
direction = self.direction.val
valid_directions = {"forward", "reverse", "bidirectional"}
if direction not in valid_directions:
raise ValueError(
"Direction {} not supported. Supported directions: {}").format(
direction, valid_directions)
weight_shape_check(self.weight_ih, self.weight_hh)
if direction == "bidirectional":
weight_shape_check(self.weight_ih_back, self.weight_hh_back)
hidden_dim, hidden_size = self.weight_hh.shape
dim_factor = 8 if direction == "bidirectional" else 4
out_seq_len = sequence_length if self.output_sequence.val else 1
num_directions = dim_factor // 4
output_shape = [out_seq_len, batch_size, num_directions * hidden_size]
output_h_shape = [batch_size, num_directions * hidden_size]
output_c_shape = [batch_size, num_directions * hidden_size]
return (
types.tensor(self.x.dtype, tuple(output_shape)),
types.tensor(self.x.dtype, tuple(output_h_shape)),
types.tensor(self.x.dtype, tuple(output_c_shape)),
)
class gru(Operation):
r"""
Gated recurrent unit (GRU).
.. math::
r_t = \rm{recurrent\_activation}(W_{ir} x_t + b_{ir} + W_{hr} h_{t-1} + b_{hr})
.. math::
z_t = \rm{recurrent\_activation}(W_{iz} x_t + b_{iz} + W_{hz} h_(t−1) + b_{hz})
.. math::
o_t = activation(W_{io} x_t + b_{io} + r_t * (W_{ho} h_(t−1) + b_{ho}))
.. math::
h_t = (1 − z_t) * o_t + z_t * h_{(t−1)}
Where:
* ``W_{ir}``, ``W_{io}``, and ``W_{iz}`` state input-hidden weight for reset, output
and update gate, respectively.
* ``W_{h[r|o|z]}`` are recurrent weights on hidden state to reset, output, update gate.
* ``h_t`` is the hidden state at time ``t``.
* ``x_t`` is the input at time ``t``.
* ``h_(t-1)`` is the hidden state of the layer at time ``t-1`` or the initial
hidden state at time ``0``.
* ``r_t``, ``o_t``, and ``z_t`` are the reset, new, and update gates, respectively.
* ``*`` is elementwise product.
Parameters
----------
x: <s, b, I, T> (Required)
* ``s`` is the sequence length, ``b`` is the batch size, and ``I`` is the
input dimension.
initial_h: <b, H, T> (Required)
* ``H`` denotes hidden size.
weight_ih: const<3*H, I, T> (Required) - Weight matrix
* ``weigh_ih = [W_{ir} | W_{io} | W_{iz}]`` where ``[a|b]`` denotes column
concatenation and ``[a, b]`` denotes row concatenation. ``W_{ir}``,
``W_{io}``, and ``W_{iz}`` have shape ``(H, I)``.
* This is used when direction="forward" or "reverse".
weight_hh: const<3*H, H, T> (Required) - Weight matrix
* ``weight_hh = [W_{hr} | W_{ho} | W_{hz}]``: ``W_{hr}``, ``W_{ho}``, and
``W_{hz}`` have shape ``(H, H)``.
* This is used when direction="forward" or "reverse".
bias: const<3*H, T> (Optional) [Default all 0s]
* ``bias[0]`` are input-hidden and hidden-hidden bias.
* ``3*H`` are biases for ``[b_{ir} + b_{hz}, b_{io}]``.
* This is used when direction="forward" or "reverse".
direction: const<str> (Optional) [Default=forward]
* Either ``forward`` or ``reverse``.
output_sequence: const<bool> (Optional) [Default=False]
* Outputs every step if ``True``.
recurrent_activation: const<str> (Optional) [Default=sigmoid]
* Activation applied on update and reset gate.
activation: const<str> (Optional) [Default=tanh]
* Activation applied on output gate.
Returns
-------
<s, b, H, T> or <1, b, H, T>
* If ``output_sequence == True`` (hidden states from every step):
``<s, b, H, T>``.
* Else ``<1, b, H, T>`` (hidden states of the final step).
<b, H, T>
* Hidden states of the final step.
Attributes
----------
T: fp32
"""
input_spec = InputSpec(x=TensorInputType(),
initial_h=TensorInputType(),
weight_ih=TensorInputType(const=True),
weight_hh=TensorInputType(const=True),
bias=TensorInputType(const=True, optional=True),
direction=StringInputType(const=True,
optional=True),
output_sequence=BoolInputType(const=True,
optional=True),
recurrent_activation=StringInputType(const=True,
optional=True),
activation=StringInputType(const=True,
optional=True))
def default_inputs(self):
return DefaultInputs(
bias=None,
direction="forward",
output_sequence=False,
recurrent_activation="sigmoid",
activation="tanh",
)
def __init__(self, **kwargs):
super(gru, self).__init__(**kwargs)
def type_inference(self):
if self.x.rank != 3:
raise ValueError(
"Invalid input shape. Expecting Rank 3 input, got {}".format(
len(self.x.rank)))
sequence_length, batch_size, input_size = self.x.shape
if self.weight_ih.rank != 2:
raise ValueError(
"Invalid weight shape. Expecting Rank 2 input, got {}".format(
len(self.weight_ih.rank)))
if self.weight_hh.rank != 2:
raise ValueError(
"Invalid weight shape. Expecting Rank 2 input, got {}".format(
len(self.weight_hh.rank)))
hidden_dim, hidden_size = self.weight_hh.shape
direction = self.direction.val
valid_directions = {"forward", "reverse"}
if direction not in valid_directions:
raise ValueError(
"Direction {} not supported. Supported directions: {}".format(
direction, valid_directions))
dim_factor = 3
if hidden_size != (hidden_dim // dim_factor):
raise ValueError(
"Incorrect weight matrix: hidden dim size mismatch. \
Provided weight_ih {}, weight_hh {}. Expecting <b, 3*H>"
).format(self.weight_ih.shape, self.weight_hh.shape)
out_seq_len = sequence_length if self.output_sequence.val else 1
output_shape = [out_seq_len, batch_size, hidden_size]
output_h_shape = [batch_size, hidden_size]
return (
types.tensor(self.x.dtype, tuple(output_shape)),
types.tensor(self.x.dtype, tuple(output_h_shape)),
)
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 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 pad(Operation):
"""
Pad a tensor.
Parameters
----------
x: tensor<[\*D_in],T> (Required)
pad: tensor<[2\*N],i32> (Required)
``N <= D_in``. Last ``N`` dimensions of ``x`` are padded as follows:
* For each dimension ``i`` of ``x`` if ``i >= D_in - N``:
* pad ``pad[2*i]`` elements before ``x[..,i,..]``
* pad ``pad[2*i+1]`` elements after ``x[..,i,..]``
* If mode is "reflect" then ``pad[2*i]`` and ``pad[2*i+1]`` can be at
most ``D[i]-1``.
* If mode is "replicate" then ``pad[2*i]`` and ``pad[2*i+1]`` can be
at most ``D[i]``.
mode: const<str> (Optional)
* Default to ``constant``.
* Must be one of the following values:
``constant``, ``reflect``, or ``replicate``.
constant_val: const<T> (Optional)
* Default to ``0``.
* Constant value to pad. Ignored if ``mode != constant``.
Returns
-------
tensor<[\*D_out],T>
* Tensor with same type as the input.
Attributes
----------
T: fp32
"""
input_spec = InputSpec(
x=TensorInputType(),
pad=IntTensorInputType(),
mode=StringInputType(const=True, optional=True),
constant_val=FloatInputType(const=True, optional=True),
)
def default_inputs(self):
return DefaultInputs(
mode="constant",
constant_val=0.,
)
def __init__(self, **kwargs):
super(pad, self).__init__(**kwargs)
def type_inference(self):
in_shape = self.x.shape
ret_shape = list(in_shape)
pad = self.pad
if len(pad.shape) != 1:
raise ValueError("Pad should be a 1D tensor!")
if self.mode and not self.mode.val in {'constant', 'reflect', 'replicate'}:
raise ValueError("Pad mode should be one of {'constant', 'reflect', 'replicate'}")
if pad.val is None:
for i in range(self.pad.shape[0]//2):
ret_shape[-self.pad.shape[0]//2+i] = get_new_symbol()
else:
pad = pad.val
pad = pad.copy()
if len(pad) % 2 != 0:
raise ValueError("Number of elements in the argument Pad must be divisible by 2.")
pad = pad.reshape(-1, 2)
if pad.shape[0] > len(ret_shape):
raise ValueError("Number of dimensions specified through pad must less than or equal to rank of input x")
for i in range(len(pad)):
ret_shape[-len(pad) + i] = ret_shape[-len(pad) + i] + pad[i][0] + pad[i][1]
return types.tensor(self.x.dtype, tuple(ret_shape))
@precondition(allow=VALUE)
def value_inference(self):
# NumPy `edge` mode is equivalent to `replicate` mode of PyTorch and CoreML
mode = "edge" if self.mode.val == "replicate" else self.mode.val
pad_val = self.pad.val
if pad_val is None:
return None
if len(self.x.val.shape) > (pad_val.shape[0] // 2):
updated_pad = np.zeros(len(self.x.val.shape) * 2)
updated_pad[-pad_val.shape[0] :] = pad_val
pad_val = updated_pad
pad_val = pad_val.reshape(-1, 2).astype(np.int32)
if mode == "constant":
return np.pad(
self.x.val, pad_val, mode, constant_values=self.constant_val.val
)
# NumPy does not support non-constant mode and constant_values argument
return np.pad(self.x.val, pad_val, mode)
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``.
Returns
-------
tensor<\*D, T>
* With the same type and shape as input ``x``.
Attributes
----------
T: fp32
"""
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(scatter, self).__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:])
np.testing.assert_equal(self.updates.shape,
np.array(expected_updates_shape))
return self.data.sym_type
class make_list(Operation):
"""
Create a list of tensor elements. The elements should have the same shape.
The list is similar to an auto-resizing array.
Parameters
----------
init_length: <i32> (Optional)
* Initial length for the list. If ``dynamic_length`` is ``False``,
``init_length`` is the fixed length of the list throughout runtime.
* Default is ``1``.
dynamic_length: <bool> (Optional)
* Initial length for the list. If ``dynamic_length`` is ``False``,
``init_length`` is the fixed length of the list throughout runtime.
* Default is ``True``.
elem_shape: <K,i32> (Required)
* Non-symbolic 1-D tensor denoting the shape of elements.
* If not provided, the resulting ``List`` won’t have the elementary shape
info, which may cause backend errors. Remedy this with SSA passes.
dtype: const<str> (Optional)
* Element tensor’s ``dtype``.
* Default is ``fp32``.
Returns
-------
List[*]
"""
input_spec = InputSpec(
init_length=IntInputType(optional=True),
dynamic_length=BoolInputType(const=True, optional=True),
elem_shape=TupleInputType(),
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(make_list, self).__init__(**kwargs)
def type_inference(self):
builtin_dtype = types.string_to_builtin(self.dtype.val)
if builtin_dtype is None:
raise ValueError("Unsupported dtype {}".format(self.dtype.val))
# Replace string with symbol
elem_shape_sym = []
for s_var in self.elem_shape:
# s is str or int
s = s_var.val
if s is None:
msg = 'make_list elem_shape must be tuple of const. ' +\
'Tuple elem {} is not'
raise ValueError(msg.format(s_var.name))
if isinstance(s, str):
try:
symbol = get_existing_symbol(s)
except ValueError:
# Must be a new symbol
symbol = get_new_symbol(s)
elem_shape_sym.append(symbol)
else:
elem_shape_sym.append(s)
elem_type = types.tensor(builtin_dtype, elem_shape_sym)
return types.list(
elem_type,
init_length=self.init_length.val,
dynamic_length=self.dynamic_length.val,
)