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, )