class select(Operation): """ Return the elements selected from either ``a`` or ``b`` depending on the ``cond``. The shape of ``cond``, ``a``, and ``b`` must be broadcastable. You must provide ``a`` and ``b`` together, or provide neither. If you provide neither, the operation returns the indices of ``cond`` that are ``True``. Parameters ---------- cond: tensor<[\*D1], B> (Required) * Tensor. When ``True``, select element from ``x``, otherwise, ``y``. a: tensor<[\*D2], T> (Optional) * Values selected at indices where ``cond`` is ``True``. * Default is ``None``. b: tensor<[\*D3], T> (Optional) * Values selected at indices where ``cond`` is ``False``. * Default is ``None``. Returns ------- tensor<[\*D_out], T> or tensor<[n, len(D1)], int32> * If ``a, b`` are both provided, the return shape is based on broadcast rules from ``cond, a, b``. * If ``a, b`` are ``None``, the return shape is 2-D, where the first dimension ``n`` is the number of matching indices in ``cond``, and ``len(D1)`` is the rank of ``cond``. Attributes ---------- B: bool T: fp32 """ input_spec = InputSpec(cond=BoolTensorInputType(), a=TensorInputType(), b=TensorInputType()) def __init__(self, **kwargs): super(select, self).__init__(**kwargs) def type_inference(self): a_type = self.a.sym_type b_type = self.b.sym_type if all([a_type, b_type]): compatible, ret_type = types.is_tensor_and_is_compatible_general_shape( a_type, b_type) if compatible: return ret_type elif a_type == b_type: return a_type else: raise ValueError("Type mismatch {} vs. {}".format( a_type, b_type)) return a_type if a_type is not None else b_type @precondition(allow=VALUE) def value_inference(self): return np.where(self.cond.val, self.a.val, self.b.val)
class conv_transpose(Operation): """ Perform transposed convolution (also known as deconvolution and fractionally stride convolution) over input. ``conv_transpose`` can also be used to compute the gradient of conv. Currently supports only 1-D and 2-D. Parameters ---------- x: tensor<[n,C_in,*D_in],T> (Required) * Input data. * ``D_in`` are spatial dimensions. * ``1 <= len(D_in) <= 2``. * ``C_in`` is the number of input channels. weight: const tensor<[C_in,C_out/groups,*D_in], T> (Required) * Filter weights. ``C_in, C_out`` are the number of input and output channels respectively. * ``D_in`` are spatial dimensions. ``1 <= len(D_in) <= 2``. bias: const tensor<[C_out],T> (Optional, default to all 0) * Bias added along output channels. pad: const tensor<[P],i32> (Optional, default to all 0s) * Number of elements to pad before and after each dimension. * ``P == 2 * len(D_in)``. * ``pad[2*i], pad[2*i+1]`` are pad sizes before and after dimension ``i``, where ``0 <= i < len(D_in)``. output_shape: const tensor<[P],i32> (Optional, default None) * Expected output shape. The first two dim must be ``[n, C_out]``. * The output shape of conv_transpose is underdetermined in general, * because conv can map multiple input shape to a single output shape. For example, for 'same' padding mode, `conv_out = ceil(conv_in/stride)`. Hence we need `output_shape` when this occurs. pad_type: const tensor<[P],i32> (Optional, default valid) * One of ``same``, ``valid``, or ``custom``. strides: const tensor<[S],i32> (Optional. Default to all 1s) * Stride along each of the spatial dimensions. * ``S == len(D_in)``. dilations: const tensor<[S],i32> (Optional. Default to all 1s) * Dilation value along each spatial dimension in ``d_in``. See ``conv``. * ``S == len(D_in)``. groups: const tensor<[], i32> (Optional. Default to 1) * Input and output channels are separated into ``groups``. * ``C_in`` and ``C_out`` must be divisible by the number of groups. See ``conv`` for examples. Returns ------- tensor<[n,C_out,*D_out],T> * ``D_out[i] = (D_in[i]-1) * strides[i] + pad[2*i] + pad[2*i+1] - (K[i] - 1) * dilations[i] + 1)] for i = 0, 1``. Attributes ---------- T: fp32 See Also -------- conv """ input_spec = InputSpec( x=FloatTensorInputType(), # [n, C_in, spatial_dims] weight=FloatTensorInputType(const=True), # [C_out, C_in, spatial_dims] bias=FloatTensorInputType(const=True, optional=True), pad=IntTensorInputType(const=True, optional=True), output_shape=IntTensorInputType(const=True, optional=True), pad_type=StringInputType(const=True, optional=True), strides=TensorInputType(const=True, optional=True), dilations=TensorInputType(const=True, optional=True), groups=IntInputType(const=True, optional=True), ) def default_inputs(self): num_spatial_dims = self.x.rank - 2 return DefaultInputs( bias=None, pad=[0]*2*num_spatial_dims, output_shape=None, pad_type="valid", strides=[1]*num_spatial_dims, dilations=[1]*num_spatial_dims, groups=1, ) def __init__(self, **kwargs): super(conv_transpose, self).__init__(**kwargs) def type_inference(self): # Input shape is [n, C_in, spatial_dims] in_shape = self.x.shape # Weight shape is [C_in, C_out/group, spatial_dims] f_shape = self.weight.shape kernel_shape = f_shape[2:] spatial_dim_rank = len(in_shape) - 2 N = in_shape[0] C_in = self.x.shape[0] groups = self.groups.val C_out = f_shape[1] * groups if self.bias is not None and self.bias.val.shape[0] != C_out: msg = "# of bias values {} not equal to # output channels {}" raise ValueError(msg.format(self.bias.val.shape[0], C_out)) if C_out % groups != 0: msg = "# of input channels {} not divisible by groups {}" raise ValueError(msg.format(C_in, groups)) # If output shape is given, return it if self.output_shape is not None: output_shape = self.output_shape.val assert output_shape[0] == N assert output_shape[1] == C_out return types.tensor( self.x.dtype, tuple(output_shape) ) strides = self.strides.val dilations = self.dilations.val kernel_shape = [ (kernel_shape[r] - 1) * dilations[r] + 1 for r in range(spatial_dim_rank) ] D_in = in_shape[2:] # spatial dimensions # Deconv's output shape is non-deterministic, we follow TF shape logic here. if self.pad_type.val == "same": d_out_shape = [strides[r] * D_in[r] for r in range(spatial_dim_rank)] elif self.pad_type.val == "valid": d_out_shape = [ strides[r] * (D_in[r]-1) + kernel_shape[r] for r in range(spatial_dim_rank) ] elif self.pad_type.val == "custom": if self.pad is None: raise ValueError("self.pad must exist if pad_type is custom") pad = self.pad.val d_out_shape = [ strides[r] * (D_in[r] - 1) + kernel_shape[r] - pad[2 * r] - pad[2 * r + 1] for r in range(spatial_dim_rank) ] retshape = [N, C_out] + d_out_shape return types.tensor(self.x.dtype, tuple(retshape))
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 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: fp32 """ 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(gather_along_axis, self).__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 scatter_along_axis(Operation): """ Scatter ``updates`` to ``data`` at locations ``indices`` at dimension ``axis`` by operation ``mode``. 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 ---------- 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_along_axis, self).__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 slice_by_index(Operation): """ Method for numpy style indexing and slicing. With a tensor ``x``, this method achieves the following: ``result = x[begin[0]: end[0]: stride[0], begin[1]: end[1]: stride[1], ...]`` Note: This method does not support pure indexing. You would need to do a squeeze if indexing is intended. Parameters ---------- x: tensor<*?, T> (Required) * Input tensor begin: tensor<[rank<x>], i32> (Required) * Starting index for the dimension of slicing. end: tensor<[rank(x)], i32> (Required) * Ending index for the dimension of slicing. stride: tensor<[rank(x)], i32> (Optional) * Default is all ``1``. * Stride for the dimension of slicing. begin_mask: tensor<[rank(x)], bool> (Optional) * Default to all ``False``. * If ``begin_mask[i]==True``, neglect ``begin[i]``, and set ``begin[i]`` to ``0``. end_mask: tensor<[rank(x)], bool> (Optional) * Default to all ``False``. * If ``end_mask[i]==True``, neglect ``end[i]``, and set ``end[i]`` to ``x.shape[i]``. squeeze_mask: tensor<[rank(x)], bool> (Optional) * Default to all ``False``. * If ``squeeze_mask[i]==true``, neglect ``end[i]``, and do the pure index at ``begin[i]``. Returns ------- tensor<\*?, T> - Scalar or tensor. Attributes ---------- T: fp32 """ input_spec = InputSpec( x=TensorInputType(), begin=IntTensorInputType(), end=IntTensorInputType(), stride=IntTensorInputType(const=True, optional=True), begin_mask=BoolTensorInputType(const=True, optional=True), end_mask=BoolTensorInputType(const=True, optional=True), squeeze_mask=BoolTensorInputType(const=True, optional=True), ) def default_inputs(self): return DefaultInputs( stride=None, begin_mask=None, end_mask=None, squeeze_mask=None, ) def __init__(self, **kwargs): super(slice_by_index, self).__init__(**kwargs) def type_inference(self): # get tensor and set default value begin = self.begin.val end = self.end.val x_rank = self.x.rank stride = self.stride.val if self.stride is not None else [1] * x_rank begin_mask = ( self.begin_mask.val if self.begin_mask is not None else [False] * x_rank ) end_mask = self.end_mask.val if self.end_mask is not None else [False] * x_rank squeeze_mask = ( self.squeeze_mask.val if self.squeeze_mask is not None else [False] * x_rank ) # solve shape x_shape = self.x.shape ret_shape = _solve_slice_by_index_shape(x_shape, begin, end, stride, begin_mask, end_mask, squeeze_mask) if len(ret_shape) == 0: # Scalar case. return self.x.dtype else: return types.tensor(self.x.dtype, tuple(ret_shape)) def value_inference(self): if self.x.sym_val is None or self.begin.val is None or self.end.val is None: return None begin = [int(i) for i in list(self.begin.val[:])] end = [int(i) for i in list(self.end.val[:])] stride = [1] * self.x.rank if self.stride is None else self.stride.val begin_mask = ( [False] * self.x.rank if self.begin_mask is None else self.begin_mask.val ) end_mask = [False] * self.x.rank if self.end_mask is None else self.end_mask.val squeeze_mask = ( [False] * self.x.rank if self.squeeze_mask is None else self.squeeze_mask.val ) slices = [] for idx, mask in enumerate(begin_mask): if mask: begin[idx] = None for idx, mask in enumerate(end_mask): if mask: end[idx] = None squeeze_axes = [] for idx, mask in enumerate(squeeze_mask): if mask: end[idx] = None stride[ idx ] = 2147483647 # We slice out only 1 element by setting stride to INF squeeze_axes.append(idx) for idx in range(self.x.rank): slices.append(slice(begin[idx], end[idx], stride[idx])) slices = tuple(slices) res = self.x.sym_val[slices] # remove squeezed axes if len(squeeze_axes) > 0: if len(squeeze_axes) == len(res.shape): if len(res) == 0: logging.warning("%s seems to be a 0 sized tensor", self.name) return np.array([]) res = res.tolist()[0] if is_symbolic(res): return res elif self.x.dtype == types.int32 or self.x.dtype == types.int64: res = np.int32(res) elif self.x.dtype == types.float or self.x.dtype == types.double: res = np.float32(res) else: raise ValueError( "Unable to convert type {}".format(self.x.sym_val.dtype) ) else: res = np.squeeze(res, axis=tuple(squeeze_axes)) return res
class squeeze(Operation): """ Remove single-dimension dimensions in a 1-D or higher tensor. Parameters ---------- x: tensor<\*?,T> (Required) * Must be at least 1-D. axes: const<K,i32> (Optional) * Axes to squeeze out. * Default to remove all single-dimensions. Returns ------- tensor<\*(rank(x)-K),T> * Tensor with same type as input ``x`` and rank ``rank(x)-K``. Attributes ---------- T: fp32 """ input_spec = InputSpec( x=TensorInputType(), axes=IntTensorInputType(const=True, optional=True), ) def default_inputs(self): return DefaultInputs( axes=None, ) def __init__(self, **kwargs): super(squeeze, self).__init__(**kwargs) def type_inference(self): x_type = self.x.dtype x_shape = self.x.shape squeezed_shape = list(x_shape) if self.axes is None: # Squeeze all single-dim, assuming symbolic dims != 1 squeezed_shape = [s for s in squeezed_shape if s != 1] else: axes = self.axes.val axes = [axis if axis >= 0 else axis + self.x.rank for axis in axes] for i in sorted(axes)[::-1]: # descending order if len(squeezed_shape) <= i: raise ValueError( "Cannot squeeze dim {} for shape {}".format(i, squeezed_shape) ) squeezed_shape.pop(i) return types.tensor(x_type, tuple(squeezed_shape)) if len(squeezed_shape) != 0 else x_type @precondition(allow=VALUE) def value_inference(self): if self.x.val is None: return None if self.axes is None: val = np.squeeze(self.x.val) else: val = np.squeeze(self.x.val, axis=tuple(self.axes.val)) return val if val.shape != () else self.x.val[0]
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: fp16, fp32, i32, bool """ input_spec = InputSpec( values=TupleInputType(), axis=IntInputType(const=True), ) def __init__(self, **kwargs): super().__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 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: fp16, 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().__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 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: fp32 """ 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(scatter_nd, self).__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 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. 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(random_uniform, self).__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 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, )
class cond(Operation): """ Perform a conditional execution. The return types must be identical between the true and false branches. Parameters ---------- pred: tensor<[], bool> (Required) * 0-D tensor (scalar) predicate to switch between true and false branches. _true_fn: function (Required) * A Python function that executes if ``pred`` evaluates to ``True``. * It must take zero input (i.e, no input), and return one or more values whose type becomes the operation's return type. _false_fn: function (Required) * A Python function that executes if ``pred`` evaluates to ``False``. * It must take zero input (i.e. no input), and have return types that match those of the ``if`` branch. Returns ------- tuple * Python tuple of ``Variables`` from one of the branches. """ input_spec = InputSpec( pred=BoolInputType(), _true_fn=PyFunctionInputType(), _false_fn=PyFunctionInputType(), ) def __init__(self, **kwargs): super(cond, self).__init__(**kwargs) def build_nested_blocks(self): # Cond block true_block_name = self.name + "_true" with Block(name=true_block_name, outer_op=self) as true_block: true_func = self._true_fn.val true_ret_vars = true_func() if isinstance(true_ret_vars, tuple): true_ret_vars = list(true_ret_vars) if not isinstance(true_ret_vars, list): true_ret_vars = [true_ret_vars] true_block.set_outputs(true_ret_vars) self.blocks.append(true_block) false_block_name = self.name + "_false" with Block(name=false_block_name, outer_op=self) as false_block: false_func = self._false_fn.val false_ret_vars = false_func() if isinstance(false_ret_vars, tuple): false_ret_vars = list(false_ret_vars) if not isinstance(false_ret_vars, list): false_ret_vars = [false_ret_vars] false_block.set_outputs(false_ret_vars) self.blocks.append(false_block) def type_inference(self): true_ret_vars = self.blocks[0].outputs false_ret_vars = self.blocks[1].outputs # Verify true_ret_vars has the same types as false_ret_vars for i, (vt, vf) in enumerate(zip(true_ret_vars, false_ret_vars)): if not is_compatible_type(vt.sym_type, vf.sym_type): msg = ("true branch output {} type {} mismatch false branch" + " output type {}") raise ValueError( msg.format(vt.name, vt.sym_type.__type_info__(), vf.sym_type.__type_info__())) return tuple(v.sym_type for v in true_ret_vars) def value_inference(self): if self.pred.val is None: raise NotImplementedError() if self.pred.val: return [v.val for v in self.blocks[0].outputs] return [v.val for v in self.blocks[1].outputs]
class while_loop(Operation): """ Perform the body repeatedly while the condition ``cond`` is true. Parameters ---------- _cond: function (Required) * A Python function that takes ``loop_vars`` as positional arguments. * The function must return a ``bool`` ``Var``. _body: function (Required) * A Python function that takes ``loop_vars`` as positional arguments. * The function must return the same number of output vars as ``loop_vars`` with the same types. loop_vars: tuple (Required) * Python tuple of ``Variables``. Returns ------- tuple * Python tuple (same type as ``loop_vars``). """ input_spec = InputSpec( # arg name with underscore prefix won't be printed. _cond=PyFunctionInputType(), _body=PyFunctionInputType(), loop_vars=TupleInputType(), ) def __init__(self, **kwargs): super(while_loop, self).__init__(**kwargs) @staticmethod def _check_equal_value(val1, val2): if val1 is None and val2 is None: return True if val1 is None or val2 is None: return False if isinstance(val1, np.ndarray) and isinstance(val2, np.ndarray): return np.array_equal(val1, val2) return val1 == val2 @staticmethod def _clean_up_child_ops(block): for op in list(block.operations): for b in op.blocks: while_loop._clean_up_child_ops(b) inputs = op.get_flattened_inputs() for in_var in inputs: in_var.remove_child_op(op) def _build_block(self, block_inputs): # Cond block: block_name = self.name + '_cond_block' with Block(block_inputs=block_inputs, outer_op=self, name=block_name) as cond_block: cond_func = self._cond.val cond_var = cond_func(*cond_block.inputs) cond_vars = cond_var if isinstance(cond_var, list) else [cond_var] cond_block.set_outputs(cond_vars) # Body block block_name = self.name + '_body_block' with Block(block_inputs=block_inputs, outer_op=self, name=block_name) as body_block: body_func = self._body.val exit_vars = body_func(*body_block.inputs) exit_vars = list(exit_vars) if isinstance(exit_vars, (list, tuple)) \ else [exit_vars] body_block.set_outputs(exit_vars) return cond_block, body_block, exit_vars def build_nested_blocks(self): # self.loop_vars is python tuple of Vars. # block_inputs Var are not produced by any op. # We assume block_inputs have the same types as self.loop_var. If not # (e.g., when certain dimensions change shape during iterate), we'd # adjust later. # We assume that sym_val is unchanging across the block iterate. If it # changes, we rebuild the block and rerun type and value inference. # Design notes on two blocks (cond and body): # # - Observe that two blocks can always be represented as a single # block that contains both cond and body logic, which would return # [loop_cond] + loop_carries. `loop_cond` is a bool. # # - Observe that single block implies a do-while logic, # in which the first iterate is always executed. It's possible to add # a cond input to while_loop to modify do-while behavior: # # %first_cond = cond_logic(...) # while_loop(cond=%first_cond, loop_vars=(...)) # # and we enter the first iterate only if cond is True. But this would # require caller to execute cond logic outside of while_loop first # (which also needs to be duplicated within the loop), # resulting in duplicated code / ops. # # - Thus, single block is unnatural for the natural execution order, # in which we execute the cond block first to get the loop_cond. Only # if `loop_cond` is True do we execute the body block. This is the # semantics of tf.while_loop. block_inputs = tuple(copy.copy(v) for v in self.loop_vars) _, visible_vars = self.enclosing_block._visible_vars_in_block() name_count = {v.name: 1 for v in visible_vars} seen = set() # Avoid using same name among block inputs for v in block_inputs: v._op = None v.op_output_idx = None v._child_ops = list() # Get unique name old_v_name = v.name if v.name in name_count: v.name = v.name + "_x" + str(name_count[v.name]) name_count[old_v_name] += 1 else: v.name = v.name + "_x0" name_count[old_v_name] = 0 v._sym_val = v._sym_val v.consuming_blocks = list() seen.add(v.name) cond_block, body_block, exit_vars = self._build_block(block_inputs) # Verify exit_vars has the same types as loop_vars block_input_type_change = False for i, (v_in, v_out) in enumerate(zip(block_inputs, exit_vars)): if not is_subtype(v_out.sym_type, v_in.sym_type): compat_shape = while_loop.get_compat_shape( v_out.sym_type, v_in.sym_type) if compat_shape is None: msg = "loop_vars '{}' changes in the body of " \ "while_loop '{}':\n {} -> {}" raise ValueError( msg.format(v_in.name, self.name, v_in.sym_type, v_out.sym_type)) else: block_inputs[i]._sym_type = types.tensor( v_in.dtype, compat_shape) block_input_type_change = True if not while_loop._check_equal_value(v_out.sym_val, v_in.sym_val): block_inputs[i]._sym_val = None block_input_type_change = True if block_input_type_change: # Since we are going to build the block again, we first need to remove ops # in the block from vars's _child_ops. while_loop._clean_up_child_ops(cond_block) while_loop._clean_up_child_ops(body_block) # Rebuild our block to invoke type inference. cond_block, body_block, exit_vars = self._build_block(block_inputs) for i, (v_in, v_out) in enumerate(zip(block_inputs, exit_vars)): if not is_subtype(v_out.sym_type, v_in.sym_type): msg = 'Block output {}: {} is not a subtype of ' +\ 'block input {}: {} after factoring shape changes' raise ValueError( msg.format(v_out.name, v_out.sym_type.__name__, v_in.name, v_in.sym_type.__name__)) if not while_loop._check_equal_value(v_out.sym_val, v_in.sym_val): msg = 'Block output {}: {} is not equal to ' +\ 'block input {}: {} after value changes' raise ValueError( msg.format(v_out.name.v.sym_val, v_in.name, v_in.sym_val)) self.blocks.append(cond_block) self.blocks.append(body_block) @staticmethod def get_compat_shape(type1, type2): """ For tensor types `type1`, `type2` that are of the same rank, return compat_shape (python list) where compat_shape[i] is integer iff type1 and type2 have the same integer shape on dim i. compat_shape[i] is symbolic otherwise. Return None if `type1`, `type2` have different rank or non-tensor type. """ if not types.is_tensor(type1) or not types.is_tensor(type2): return None s1 = type1.get_shape() s2 = type2.get_shape() if len(s1) != len(s2): return None compat_shape = [] for d1, d2 in zip(s1, s2): if d1 != d2: compat_shape.append(get_new_symbol()) else: compat_shape.append(d1) return compat_shape def type_inference(self): # Skip the conditional var return tuple(v.sym_type for v in self.blocks[1].outputs)
class expand_dims(Operation): """ Insert a single-dimension in a 1-D or higher tensor at each axis in axes. Parameters ---------- x: tensor<\*?, T> (Required) * Scalar or tensor. axes: const tensor<[K], i32> Required * ``K`` is the number of dimensions expanded. * Insert single dimension at dimension index at each axes. * Negative value to index from the end. ``-d-1 <= axis <= d`` where ``d`` is the rank of ``x``. Returns ------- tensor<\*(rank(x)+K), T> * Same type as the input ``x`` with rank ``rank(x)+K``. Attributes ---------- T: fp32 """ input_spec = InputSpec( x=ScalarOrTensorInputType(), axes=IntTensorInputType(const=True), ) def __init__(self, **kwargs): super(expand_dims, self).__init__(**kwargs) def type_inference(self): x_rank = self.x.rank x_type = self.x.dtype x_shape = list(self.x.shape) axes = self.axes.val out_rank = x_rank + len(axes) for axis in axes: if axis <= -out_rank - 1 or axis >= out_rank: msg = 'Axis value {} is out of bounds for {} node "{}" of shape {}' raise IndexError( msg.format(axis, self.op_type, self.name, self.x.shape) ) ret_shape = x_shape axes = sorted([out_rank + axis if axis < 0 else axis for axis in axes]) for axis in axes: ret_shape.insert(axis, 1) return types.tensor(x_type, tuple(ret_shape)) @precondition(allow=VALUE) def value_inference(self): axes = self.axes.val out_rank = self.x.rank + len(axes) for axis in axes: if axis <= -out_rank - 1 or axis >= out_rank: msg = 'Axis value {} is out of bounds for {} node "{}" of shape {}' raise IndexError( msg.format(axis, self.op_type, self.name, self.x.shape) ) axes = sorted([out_rank + axis if axis < 0 else axis for axis in axes]) ret_shape = list(self.x.shape) for axis in axes: ret_shape.insert(axis, 1) return np.reshape(self.x.val, ret_shape)
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: fp16, fp32, i32, bool """ 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().__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)
class reshape(Operation): """ Return a tensor that has the same values as ``x`` with shape ``shape``. ``shape`` must have the same volume (number of elements) as ``x``. Parameters ---------- x: tensor<\*?, T> (Required) * A n-D tensor or a scalar. * If ``x`` is fixed rank (and possibly contains symbolic dimension), shape may contain elements that are not positive integers (see below). * If ``x`` is variadic rank, shape can only contain positive integers. shape: tensor<[K], i32> (Required) A 1-D tensor, with elements from the following: * Positive integers. * Symbols: All but one symbol in shape must be present in ``x.shape``. The new symbol that is not present in ``x.shape`` represent a dimension such that the total size remains constant. Symbol is illegal if ``x`` is variadic rank. * ``-1``: ``-1`` introduces a new symbol (see Symbols). Therefore, ``-1`` is allowed if all symbols in the shape appear in ``x.shape``. ``-1`` is illegal if ``x`` is variadic rank. * ``0``: If ``K == rank(x)`` then ``0`` means inheriting from the corresponding dimension in ``x.shape``. ``0`` is illegal if ``x`` is variadic rank. Returns ------- tensor<\*?, T> * Tensor with shape determined by the input shape. Attributes ---------- T: fp32 """ input_spec = InputSpec( x=ScalarOrTensorInputType(), shape=IntTensorInputType(), ) def __init__(self, **kwargs): super(reshape, self).__init__(**kwargs) def type_inference(self): if any_symbolic(self.shape.shape): # We can't infer any shape if shape has variable length. return types.tensor(self.x.dtype, (get_new_variadic_symbol(),)) # shape has fixed length here. if self.shape.sym_val is None: shape = tuple([get_new_symbol() for _ in range(self.shape.shape[0])]) return types.tensor(self.x.dtype, shape) t, _ = self._get_type_val() return t @precondition(allow=VALUE | SYMBOL) def value_inference(self): _, val = self._get_type_val() return val def _get_type_val(self): x_type = self.x.dtype x_shape = self.x.shape x_vol = np.prod(x_shape) # shape is const, and thus sym_val is not None sym_shape = self.shape.sym_val sym_shape = [get_new_symbol() if d == -1 else d for d in sym_shape] try: ret_shape = reshape.enforce_volumetric_constraint(x_vol, sym_shape) except: ret_shape = sym_shape ret_val = None if self.x.val is not None and all(isscalar(a) and not is_symbolic(a) for a in ret_shape): ret_val = reshape_with_symbol(self.x.val, ret_shape) return types.tensor(x_type, tuple(ret_shape)), ret_val @staticmethod def enforce_volumetric_constraint(left_volume, inshape): left_symbols = set() if is_symbolic(left_volume): left_symbols = left_volume.free_symbols # Generally, we want to solve for right in terms of left. But this # is kinda annoying actually. shape = list(inshape) # Handling when reshape is given 0 instead of actual input # input tensor shape: [4, 3, 2], reshape:[0, -1], output tensor shape: [4, 6] if shape.count(-1) > 1: raise ValueError( "Reshape op supports only one dimension to be -1. Given {}".format( shape.count(-1) ) ) infer_dim_index = shape.index(-1) if -1 in shape else None right_volume = 1 for i in shape: if i != -1: right_volume = right_volume * i if infer_dim_index: shape[infer_dim_index] = left_volume // right_volume if not is_symbolic(right_volume): return shape constraints = [left_volume - right_volume] solve_for = [s for s in shape if is_symbolic(s)] for rightsym in solve_for: sol = sm.solve(constraints, [rightsym], dict=True) if not isinstance(sol, list): sol = [sol] # look for an acceptable solution for s in sol: if 0 in s.values(): continue for i in range(len(shape)): if shape[i] in s: v = s[shape[i]] if len(v.free_symbols - left_symbols) > 0: continue try: shape[i] = int(v) except: shape[i] = v return shape
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: fp16, 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().__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 slice_by_size(Operation): """ Slice input tensor starting from the given ``begin`` index and by the amount specified by the ``size`` input, for each dimension. Parameters ---------- x: tensor<*?, T> (Required) * Input tensor. begin: tensor<[rank(x)], i32> Required * The begin index for slice. size: tensor<[rank(x)], i32> Required * The size that is to be sliced. If ``size`` is ``-1``, all the remaining elements starting with "begin" are sliced. Returns ------- tensor<\*?, T> * Scalar or tensor. Attributes ---------- T: fp32 """ input_spec = InputSpec( x=TensorInputType(), begin=IntTensorInputType(), size=IntTensorInputType(), ) def __init__(self, **kwargs): super(slice_by_size, self).__init__(**kwargs) def type_inference(self): if self.begin.rank != 1: raise ValueError( "begin should be 1-D tensor, got {}-D tensor instead".format( self.begin.rank ) ) if self.size.rank != 1: raise ValueError( "size should be 1-D tensor, got {}-D tensor instead".format( self.size.rank ) ) if self.x.rank != self.begin.shape[0]: raise ValueError( "Length of begin {} doesn't equal to input rank {}.".format( len(self.begin.shape[0]), len(self.x.rank) ) ) if self.x.rank != self.size.shape[0]: raise ValueError( "Length of size {} doesn't equal to input rank {}.".format( len(self.size.shape[0]), len(self.x.rank) ) ) x_shape = self.x.shape ret_shape = [] if self.size.sym_val is None: ret_shape = [get_new_symbol() for _ in range(self.x.rank)] return types.tensor(self.x.dtype, tuple(ret_shape)) for idx, s in enumerate(self.size.sym_val): if is_symbolic(s): ret_shape.append(s) elif s != -1: ret_shape.append(s) elif self.begin.sym_val is not None: ret_shape.append(x_shape[idx] - self.begin.sym_val[idx]) else: ret_shape.append(get_new_symbol()) return types.tensor(self.x.dtype, tuple(ret_shape)) @precondition(allow=VALUE | SYMBOL) def value_inference(self): if any_symbolic(self.begin.sym_val): return None if any_symbolic(self.size.sym_val): return None if self.x.val is None: return None slices = [] for i in range(self.x.rank): begin_val = self.begin.val[i] if begin_val < 0: if is_symbolic(self.x.shape[i]): return None begin_val += self.x.shape[i] if self.size.val[i] > 0: slices.append(slice(begin_val, begin_val + self.size.val[i])) else: slices.append(slice(begin_val, None, None)) return self.x.val[tuple(slices)]
class range_1d(Operation): """ Returns a numpy-like 1- range sequence. Parameters ---------- end: <T> (Required) * The upper limit of the sequence, exclusive. start: <T> (Required) * The start point of the sequence. step: <T> (Required) * Number that increments ``start``. Returns ------- tensor<M, T> * A 1-D tensor, where ``M`` is the length of the sequence. Attributes ---------- T: i32, fp16, fp32 """ input_spec = InputSpec( end=IntOrFloatInputType(), start=IntOrFloatInputType(), step=IntOrFloatInputType(), ) def __init__(self, **kwargs): super().__init__(**kwargs) @precondition(allow=VALUE) def value_inference(self): start = self.start.val end = self.end.val step = self.step.val shape = (end - start) / step # To prevent from creating constant greater then 1MB, # a upper bound of the size of the resulting array is set. if shape > MAX_SIZE_CONSTANT_FOLDING: return None return np.arange(start, end, step) def type_inference(self): start = self.start.sym_val end = self.end.sym_val step = self.step.sym_val if ((self.start.dtype != self.end.dtype) or (self.start.dtype != self.step.dtype) or (self.end.dtype != self.step.dtype)): raise TypeError( "All inputs to the range operation must have same input types." ) if all(sym_val is not None for sym_val in (start, end, step)): shape = (end - start) / step shape = shape if is_symbolic(shape) else int(math.ceil(shape)) shape = tuple([shape]) else: shape = tuple([ get_new_symbol(), ]) return types.tensor(self.start.dtype, 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 """ 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:]) 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 tile(Operation): """ Returns a new tensor by replicating input ``x`` multiples times. Dimension ``i`` of ``x`` will be replicated ``reps[i]`` times. Parameters ---------- x: tensor<\*?, T> (Required) * Input tensor. reps: tensor<[rank(x)], int32> (Required) * A 1-D tensor with length ``rank(x)``, which indicates the number to replicate the input along each dimension. Returns ------- tensor<\*?, T>: * An n-D tensor with same type as the input. Attributes ---------- T: fp16, fp32, i32, bool """ input_spec = InputSpec( x=TensorInputType(), reps=TensorInputType(), ) def __init__(self, **kwargs): super().__init__(**kwargs) def type_inference(self): x_type = self.x.dtype x_shape = np.array(self.x.shape) reps = self.reps.sym_val if reps is None: out_shape = tuple([get_new_symbol() for _ in range(self.x.rank)]) return types.tensor(x_type, out_shape) if len(reps) == 0 or len(reps) != self.x.rank: msg = ("Length of the reps ({}) must be at least 1, and " "equal to the rank of the input x ({})") raise ValueError(msg.format(len(reps), self.x.rank)) out_shape = [] for i, rep in enumerate(reps): if not is_symbolic(rep): if rep <= 0: raise ValueError( "All entries of reps parameter must be greater than 0") if is_symbolic(rep) or is_symbolic(x_shape[i]): out_shape.append(get_new_symbol()) else: out_shape.append(rep * x_shape[i]) out_shape = tuple(out_shape) return types.tensor(x_type, out_shape) @precondition(allow=VALUE) def value_inference(self): # Infer only if don't have symbolic values. if self.reps.val is None: return None return np.tile(self.x.val, reps=self.reps.val)
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 ---------- T: fp32 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(gather, self).__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 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: fp16, fp32, i32 """ 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().__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 conv(Operation): """ Perform convolution over input. Currently supports only 1-D and 2-D convolution. Parameters ---------- x: tensor<[n, C_in, \*d_in], T> (Required) * ``d_in`` are (possibly runtime-determined) spatial dimensions. For example, ``d_in = [224, 224]`` for 2D convolution. * ``1 <= len(d_in) <= 2``: Only 1-D and 2-D convolution. * ``C_in`` is the number of input channels or depth dimensions. * ``n`` is the batch dimension. weight: tensor<[C_out, C_in/groups, \*K], T> (Required) * Filter weights. * ``C_in`` is the number of input channels. * ``C_in`` must be divisible by ``groups``. * ``K`` are kernel sizes. For example, ``K = [KH, KW]`` for 2-D conv. * When ``dilations`` is not all ``1``, ``weight`` has to be ``const`` at compile time strides: const tensor<[S], i32> (Optional) * Default to one vector of length equal to the number of spatial dimensions. * Strides along each of the spatial dimensions. * ``S == len(d_in)``. pad_type: const str (Required) Must be one of the following: * ``valid``: No padding. This is equivalent to custom pad with ``pad[2*i] == pad[2*i+1] == 0, for i=0,...,len(d_in)-1``. * ``custom``: Specify custom padding in the parameter ``pad``. * ``same``: input is padded such that out spatial shapes are ``d_out[i] = ceil(d_in[i] / strides[i])``. Specifically, for ``i = 0,..,,len(d_in)-1``, the equivalent paddings are calculated as follows: * ``dilated_kernel = (K[i] - 1) * dilate[i] + 1`` * if ``dilated_kernel`` is odd, ``padding[2*i] = padding[2*i+1] = floor(dilated_kernel / 2)`` * Otherwise: ``padding[2*i] = ceil((dilated_kernel - 1) / 2)``, ``padding[2*i+1] = floor((dilated_kernel - 1) / 2)`` pad: const tensor<[P], i32> (Optional. Default to all zeros) * ``len(P) = 2 * len(d_in)`` * ``pad`` should be specified if and only if ``pad_type == custom``, otherwise errors occur. * ``pad`` represents the number of elements to pad before and after each dimension. Specifically, ``pad[0], pad[1]`` are the pad size before / after spatial dimension 0, ``pad[2], pad[3]`` are the pad size before / after spatial dimension 1, etc. dilations: const tensor<[S], i32> (Optional. Default to all 1s) * Dilation value along each spatial dimension in ``d_in``. See `visualization <https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md>`_. * ``S == len(d_in)``. groups: const tensor<[], i32> (Optional, default to 1) * Input and output channels are split by ``groups``. * ``C_in`` must be divisible by ``groups``. * Maximum value for group is ``C_in``, in which case it is a depthwise convolution. For examples (assuming ``C_in = 16, C_out = 32``): * ``groups == 1``, ``weight`` has shape ``[32, 16, KH, KW]``: All input channels are convolved with the ``weight`` kernel to produce all output channels. * ``groups == 2``, ``weight`` has shape ``[32, 8, KH, KW]``: Input channels 0~7 are convolved with half of the ``weight`` kernel to produce output channels 0~15. Similarly, input channels 8~15 are convolved with the other half of ``weight`` to product output channels 16~31. * ``groups == C_in``, ``weight`` has shape ``[32, 1, KH, KW]``: Each input channel is convolved with its own set of filters and each produce ``C_out / C_in = 2`` channels. This is equivalent to depthwise convolution. bias: const tensor<[C_out],T> (Optional, default to all 0) * Bias along output channels. Returns ------- tensor<[n, C_out, \*d_out], T> * Output activation has the same rank and spatial dimension as the input. That is, ``len(d_out) == len(d_in)``. * For ``i=0,..,len(d_in)-1, d_out[i] = floor [(D_in[i] + pad[2*i] + pad[2*i+1] - (K[i]-1)*dilations[i] - 1) / strides[i] ] + 1`` Attributes ---------- T: fp32 See Also -------- conv_transpose """ input_spec = InputSpec( x=FloatTensorInputType(), weight=FloatTensorInputType(), bias=FloatTensorInputType(const=True, optional=True), strides=IntTensorInputType(const=True, optional=True), pad_type=StringInputType(const=True, optional=True), pad=IntTensorInputType(const=True, optional=True), dilations=IntTensorInputType(const=True, optional=True), groups=IntInputType(const=True, optional=True), ) def default_inputs(self): num_spatial_dims = self.x.rank - 2 return DefaultInputs( bias=None, strides=[1]*num_spatial_dims, pad_type="valid", pad=[0]*num_spatial_dims*2, dilations=[1]*num_spatial_dims, groups=1, ) def __init__(self, **kwargs): super(conv, self).__init__(**kwargs) def type_inference(self): inshape = self.x.shape f_shape = self.weight.shape kernel_shape = f_shape[2:] C_out = f_shape[0] C_in = self.x.shape[1] groups = self.groups.val if self.bias is not None and self.bias.val.shape[0] != C_out: msg = "# of bias values {} not equal to # output channels {}" if C_in % groups != 0: msg = "# of input channels {} not divisible by groups {}" raise ValueError(msg.format(C_in, groups)) if C_in // groups != self.weight.shape[1]: msg = "C_in / groups = {}/{} != weight[1] ({})" raise ValueError(msg.format(C_in, groups, self.weight.shape[1])) strides = self.strides.val dilations = self.dilations.val # Ignore self.pad if pad_type != custom custom_pad = None if self.pad_type.val != 'custom' else self.pad.val if self.weight.val is None and any([True if d > 1 else False for d in dilations]): raise ValueError("Convolution with dynamic weights does not support dilations!") N = inshape[0] C_out = f_shape[0] # spatial dimensions d_out_shape = spatial_dimensions_out_shape( pad_type=self.pad_type.val, input_shape=inshape[2:], kernel_shape=kernel_shape, strides=strides, dilations=dilations, custom_pad=custom_pad, ) retshape = [N, C_out] + d_out_shape return types.tensor(self.x.dtype, tuple(retshape))
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: fp16, fp32, i32 """ 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().__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 avg_pool(Pooling): """ Perform average pooling. Supports 1-D, 2-D, and 3-D pool (1, 2, or 3 spatial dimensions). Parameters ---------- x: tensor<[n,C_in,\*D_in],T> (Required) * ``3 <= rank <= 5``. * ``D_in`` are spatial dimensions, ``1 <= len(D_in) <= 3``. * ``C_in`` is the number of input channels or depth dimensions. * ``n`` is the batch dimension. kernel_sizes: const tensor<[K],T> (Required) * The size of the window for each spatial dimension ``D_in`` of the input tensor. * ``K == len(D_in)`` strides: const tensor<[S],i32> (Optional, default to all 1s) * Stride along each of the spatial dimensions. * ``S == len(D_in)``. pad_type: const str (Required) Must be one of ``valid``, ``same`` or ``custom``. * ``valid``: No padding. This is equivalent to custom pad with ``pad[i] = 0, for all i``. * ``same`` : This is equivalent to custom pad with ``pad[2*i] + pad[2*i+1] = kernel_size[i]``. * ``custom``: Specify custom padding in the parameter pad. note that ``same`` padding is equivalent to custom padding with ``pad[2*i] + pad[2*i+1] = kernel_size[i]``. pad: const<[P],i32> (Optional. Default to all 0s) * ``pad`` represents the number of elements to pad before and after each dimension: ``pad[2*i], pad[2*i+1]`` are the pad size before and after spatial dimension ``i``. * ``P = 2 * len(D_in)``. * ``pad`` should be specified if and only if ``pad_type == custom`` exclude_padding_from_average: const tensor<[], bool> (Optional, default to False) * If ``True``, padded values (0s) are excluded from the denominator count when computing the average over the kernel window. ceil_mode: const<bool> * Same as PyTorch's ``ceil`` mode. * ``ceil`` is used instead of floor in calculating the output size. * Optional, defaults to ``False``. * Only applicable when ``pad_type`` is ``valid`` or ``custom``. * When ``ceil_mode`` is True, padding must be symmetric; that is, if specified, ``pad[2*i] == pad[2*i+1]`` must hold. Returns ------- tensor<[n, C_out,\*D_out],T> * Same rank as ``x``. * When ``ceil_mode = False``: * ``D_out[i] = floor[(D_in[i] + pad[2*i] + pad[2*i+1] - kernel_sizes[i]) / strides[i]] +1, for i = 0, .., len(D_in) - 1`` is mathematically the same as (when all parameters involved are integers): * ``D_out[i] = ceil [(D_in[i] + pad[2*i] + pad[2*i+1] - kernel_size[i] - 1) / stride[i]], for i = 0, .., len(D_in) - 1``. * ``*D_out`` is all ones if ``global_pooling`` is ``true``. * When ``ceil_mode = True``: * ``D_out[i] = ceil[(D_in[i] + pad[2*i] + pad[2*i+1] - kernel_sizes[i]) / strides[i]] +1, for i = 0, .., len(D_in) - 1`` * If ``(D_out[i] - 1) * strides[i] >= D_in[i] + pad[2*i] and (pad[2*i] + pad[2*i+1] > 0)`` then ``D_out[i] = D_out[i] - 1``. * The first equation is same as: * ``D_out[i] = floor[(D_in[i] + pad[2*i] + pad[2*i+1] - kernel_sizes[i] + strides[i] - 1) / strides[i]] +1, for i = 0, .., len(D_in) - 1`` Attributes ---------- T: fp16, fp32 See Also -------- l2_pool, max_pool """ input_spec = ( InputSpec( exclude_padding_from_average=BoolInputType(const=True, optional=True)) + Pooling.input_spec ) def default_inputs(self): return super().default_inputs() + \ DefaultInputs( exclude_padding_from_average=False, ) def __init__(self, **kwargs): super().__init__(**kwargs)
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: fp16, fp32, i32, bool """ 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().__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 split(Operation): """ Split tensors into a tuple Parameters ---------- x: <\*?,T> (Required) * The tensor to split. * The tensors may be variadic, but the number of tensors must be determined at compile time (i.e. a tuple). num_splits: <i32> (Optional) If specified, divide ``x`` into ``num_splits`` tensors along ``axis``. Its behavior depends on ``split_sizes``: * If ``split_sizes`` is defined, ``num_splits == S``, and the output sizes may be uneven. * If ``split_sizes`` is not defined, ``value.shape[axis]`` must be divisible by ``num_splits``, and the output sizes must be even. At least one of ``num_splits`` or ``split_sizes`` must be provided. If ``split_sizes`` length ``S`` cannot be determined at compile time, ``num_splits`` must be supplied to determine the number of outputs. split_sizes: const<S,i32> (Optional) * Sizes to split to. The sum of ``split_sizes`` must equal to ``value.shape[axis]``. axis: const<i32> (Required) * The dimension along which to concatenate. Must be in the range ``[-rank(x), rank(x))``. Returns ------- Tuple[tensor<\*?,T>] * Where the length of the tuple is the number of splits (determined from ``num_splits`` or ``split_sizes``). Attributes ---------- T: fp32 """ input_spec = InputSpec( x=TensorInputType(), num_splits=IntInputType(const=True, optional=True), split_sizes=IntTensorInputType(const=True, optional=True), axis=IntInputType(const=True), ) def __init__(self, **kwargs): super(split, self).__init__(**kwargs) def type_inference(self): num_splits, sizes = self._get_num_splits_and_sizes() x_shape = list(self.x.shape) ret_shapes = [x_shape[:] for _ in range(num_splits)] axis = self.axis.val for i, d in enumerate(sizes): ret_shapes[i][axis] = d self.sizes = sizes return tuple([types.tensor(self.x.dtype, s) for s in ret_shapes]) def _get_num_splits_and_sizes(self): """ Return: - num_splits: int - sizes: list of int/symbols. Of length num_splits Raise ValueError if num_splits cannot be determined. """ if self.num_splits is None and self.split_sizes is None: msg = ( "At least one of num_splits and split_sizes " + "must be specified in split op {}" ) raise ValueError(msg.format(self.name)) axis = self.axis.val if self.num_splits is not None: num_splits = self.num_splits.val if self.split_sizes is None: # Even split if ( not is_symbolic(self.x.shape[axis]) and self.x.shape[axis] % num_splits != 0 ): msg = "num_split {} does not divide split " + "dim (length = {})" raise ValueError(msg.format(num_splits, self.x.shape[axis])) size = self.x.shape[axis] / num_splits return num_splits, [size] * num_splits # self.split_sizes is not None if self.split_sizes.sym_val is not None: return num_splits, self.split_sizes.sym_val # self.split_size.sym_val is None. sizes = [get_new_symbol() for _ in range(num_splits)] return num_splits, sizes # self.num_splits is None, self.split_sizes is not None if self.split_sizes.sym_val is not None: return len(self.split_sizes.sym_val), self.split_sizes.sym_val # self.num_splits is None, self.split_sizes is not None # self.split_sizes.sym_val is None if any_symbolic(self.split_sizes.shape): raise ValueError("Unable to determine number of splits") num_splits = len(self.split_sizes.shape) sizes = [get_new_symbol() for _ in range(num_splits)] return num_splits, sizes @precondition(allow=VALUE | SYMBOL | NONE) def value_inference(self): num_splits, sizes = self._get_num_splits_and_sizes() if self.x.sym_val is None or any_symbolic(sizes): raise NotImplementedError() if num_splits == 1: # No split_indices possible. return self.x.sym_val split_indices = np.cumsum(sizes).astype(np.int) return tuple(np.split(self.x.sym_val, split_indices[:-1], axis=self.axis.val))
class Const(Operation): """ A base class that returns constant values. Parameters ---------- mode: immediate_value, file_value (Optional) * Determines how the constant value is stored in the internal MIL format. * For large constants such as convolution weights, use ``file_value``. * For smaller-size constants such as values of a stride, use ``immediate_value``. val: const<\*,T> (Required) Returns ------- const<\*,T> Attributes ---------- T: fp32, i32, str """ input_spec = InputSpec(val=InternalScalarOrTensorInputType(const=True), ) def __init__(self, **kwargs): super(Const, self).__init__(**kwargs) def type_inference(self): builtin_type, _ = self._get_type_val(self.val.val) return builtin_type def value_inference(self): _, val = self._get_type_val(self.val.val) return val def _get_type_val(self, value): if isinstance(value, (float, np.float64)): value = np.float32(value) elif isinstance(value, bool): value = np.bool(value) elif isinstance(value, (int, np.int64)): value = np.int32(value) elif isinstance(value, (tuple, list, np.ndarray)): value = np.array(value) # For the int type, we use int32 by default if value.dtype in [ np.uint8, np.int8, np.uint16, np.int16, np.uint32, np.uint64, np.int64 ]: if value.dtype in [np.uint64, np.int64]: msg = "Downcast const op {} data int64 as int32".format( self.name) logging.debug(msg) value = value.astype(np.int32) # For the float type, we use float32 by default elif value.dtype == np.float64: msg = "Downcast const op {} data fp64 as fp32".format( self.name) logging.debug(msg) value = value.astype(np.float32) elif isinstance(value, mil_list): # if val that was passed in is of type mil_list, which is just a wrapper on top of python list # then construct the list type list_value = value.ls if len(list_value) == 0: raise ValueError("'mil_list' points to an empty list") builtin_elem_type, _ = self._get_type_val(list_value[0]) # mil_list is a special case that we want to preserve the int64 element type if isinstance(list_value[0], np.int64): builtin_elem_type = types.int64 from coremltools.converters.mil.mil.types.type_list import list as types_list builtin_type = types_list(builtin_elem_type, init_length=len(list_value), dynamic_length=False) return builtin_type, value if not isinstance(value, (np.generic, np.ndarray, str, bool, mil_list)): raise ValueError("Unknown value for constant: {}".format(value)) _, builtin_type = numpy_val_to_builtin_val(value) return builtin_type, value