Source code for coremltools.converters.mil.mil.ops.defs.tensor_operation

#  Copyright (c) 2020, Apple Inc. All rights reserved.
#
#  Use of this source code is governed by a BSD-3-clause license that can be
#  found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause

import math
import numpy as np

from coremltools.converters.mil.mil.types.symbolic import (
    is_symbolic,
    any_symbolic,
    is_compatible_symbolic_vector,
)
from coremltools.converters.mil.mil import (
    get_new_symbol,
    get_new_variadic_symbol,
    SYMBOL,
    VALUE,
    NONE,
)

from coremltools.converters.mil.mil import types
from ._op_reqs import *
from ._utils import promoted_primitive_type


[docs]@register_op(doc_str="") class band_part(Operation): """ Returns a tensor setting everything outside a center band to zeros for the innermost matrix. Special cases: - ``band_part(x, 0, -1)`` returns upper triangular part. - ``band_part(x, -1, 0)`` returns lower triangular part. - ``band_part(x, 0, 0)`` returns diagonal. Parameters ---------- x: tensor<\*?, T> (Required) * Input tensor. lower: const<i32> (Optional) * Number of lower / below sub-diagonals to keep. If negative, keep entire lower triangle. * Defaults to ``-1`` (keep the entire lower triangle). upper: const<i32> (Optional) * Number of upper / above sub-diagonals to keep. If negative, keep entire lower triangle. * Defaults to ``-1`` (keep the entire upper triangle). Returns ------- tensor<\*?, T> * Same type and shape as the input tensor. """ input_spec = InputSpec( x=TensorInputType(), lower=IntInputType(const=True, optional=True), upper=IntInputType(const=True, optional=True), ) def default_inputs(self): return DefaultInputs( lower=-1, upper=-1) def __init__(self, **kwargs): super(band_part, self).__init__(**kwargs) def type_inference(self): return self.x.sym_type
[docs]@register_op(doc_str="") class cumsum(Operation): """ Returns the cumulative sum of the input along the given axis. Parameters ---------- x: tensor<\*?, T> (Required) * Input tensor. axis: const<i32> (Optional) * default to ``0``. * Axis for which the cumulative sum is computed. exclusive: const<bool> (Optional) * Default to ``False``. * When set to ``False``, inclusive cumsum is computed, that is the first element of the output is identical to the first element in the input. * When set to ``True``, exclusive cumsum is computed, which makes the first element of output to ``0``. reverse: const<bool> (Optional) * Default to ``False``. * When set to ``True``, perform cumsum in the reverse order. Returns ------- tensor<\*?, T> * Same type and shape as the input tensor. Attributes ---------- T: fp32, int32 """ input_spec = InputSpec( x=TensorInputType(), axis=IntInputType(const=True, optional=True), exclusive=BoolInputType(const=True, optional=True), reverse=BoolInputType(const=True, optional=True), ) def default_inputs(self): return DefaultInputs( axis=0, exclusive=False, reverse=False) def __init__(self, **kwargs): super(cumsum, self).__init__(**kwargs) @precondition(allow=VALUE) def value_inference(self): data = np.copy(self.x.val) axis = self.axis.val reverse = self.reverse.val exclusive = self.exclusive.val if reverse: data = np.flip(data, axis=axis) data = np.cumsum(data, axis=axis) if exclusive: zero_shape = np.copy(data.shape) zero_shape[axis] = 1 data = np.concatenate((np.zeros(zero_shape, data)), axis=axis) if reverse: data = np.flip(data, axis=axis) return data def type_inference(self): # Check range of axis if self.axis.val < -1 or self.axis.val > self.x.rank - 1: raise ValueError( "axis should be in the range [-1, {}]".format(self.x.rank - 1) ) return self.x.sym_type
[docs]@register_op(doc_str="") class fill(Operation): """ Returns a tensor with a given shape filled with a constant value. Parameters ---------- shape: tensor<[K], i32> (Required) * Target output tensor shape. * ``K`` is the rank of the output tensor. ``shape[k] > 0`` for ``k = 0,..., K-1``. value: const<T> (Optional) * Default to ``0.0``. * Constant value to fill in. Returns ------- tensor<\*?, T> * Tensor with shape determined by the input shape. Attributes ---------- T: fp32 """ input_spec = InputSpec( shape=IntTensorInputType(), value=IntOrFloatOrBoolInputType(const=True, optional=True), ) def default_inputs(self): return DefaultInputs( value=0.) def __init__(self, **kwargs): super(fill, 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.value.dtype, (get_new_variadic_symbol(),)) # shape has fixed length here. if self.shape.sym_val is None: ret_shape = tuple([get_new_symbol() for _ in range(self.shape.shape[0])]) return types.tensor(self.value.dtype, ret_shape) return types.tensor(self.value.dtype, tuple(self.shape.sym_val.tolist())) @precondition(allow=VALUE) def value_inference(self): return np.full(shape=self.shape.val, fill_value=self.value.val)
[docs]@register_op(doc_str="") class non_maximum_suppression(Operation): """ Applies non-maximum suppression (NMS) on the input box coordinates according to their intersection-over-union (IoU). NMS selects a subset of bounding boxes in descending order of score, and removes boxes that have high intersection-over-union (IOU) overlap with previously-selected boxes. Parameters ---------- boxes: tensor<[n, B, 4], T> (Required) * Box coordinates on which to perform NMS. scores: tensor<[n, B, K], T> (Required) * Scores for each one of the boxes. iou_threshold: const<T> (Required) * The intersection over union (``IoU``) threshold over which boxes are suppressed. NMS remove all overlapping boxes with ``IoU > iou_threshold``. score_threshold: const<T> (Required) * Before IoU suppression is performed, boxes with class scores below this threshold are rejected. max_boxes: const<i32> (Required) * Maximum number of boxes to select. If the number of surviving boxes are less, output is padded up to this number. per_class_suppression: const<bool> (Optional) * Default to ``False``. * If ``True``, suppression is performed independently within boxes of each class. Returns ------- tensor<[n, max_boxes, 4], T> * Coordinates of selected boxes. tensor<[n, max_boxes, K], T> * Scores of selected boxes. tensor<[n, max_boxes], i32> * Indices of selected boxes. tensor<[n], i32> * Number of boxes selected for each batch. Attributes ---------- T: fp32 """ input_spec = InputSpec( boxes=TensorInputType(), scores=TensorInputType(), iou_threshold=FloatInputType(const=True), score_threshold=FloatInputType(const=True), max_boxes=IntInputType(const=True), per_class_suppression=BoolInputType(const=True, optional=True), ) def default_inputs(self): return DefaultInputs( per_class_suppression=False) def __init__(self, **kwargs): super(non_maximum_suppression, self).__init__(**kwargs) def type_inference(self): boxes_dtype = self.boxes.dtype scores_dtype = self.scores.dtype n_batch, _, n_score = self.scores.shape max_boxes = self.max_boxes.val return ( types.tensor(boxes_dtype, (n_batch, max_boxes, 4)), types.tensor(scores_dtype, (n_batch, max_boxes, n_score)), types.tensor(types.int32, (n_batch, max_boxes)), types.tensor(types.int32, (n_batch,)), )
[docs]@register_op(doc_str="") class non_zero(Operation): """ Returns the indices of the elements in the given tensor that are non-zero. Parameters ---------- x: tensor<\*?, T> (Required) * Tensor, values selected at indices where its values is not equal to ``0``. Returns ------- tensor<[N, R], int32> * 2-dimensional tensor contains indices of elements that are non-zero. Each row is the index for a non-zero value. * ``N`` is the number of non-zero elements, ``R`` is the rank of the input. Attributes ---------- T: fp32, int32 """ input_spec = InputSpec(x=TensorInputType()) def __init__(self, **kwargs): super(non_zero, self).__init__(**kwargs) def type_inference(self): shape = tuple([get_new_symbol(), self.x.rank]) return types.tensor(types.int32, shape) @precondition(allow=VALUE) def value_inference(self): return np.transpose(np.nonzero(self.x.val))
[docs]@register_op(doc_str="") class one_hot(Operation): """ Returns one-hot vectors whose locations represented in ``indices`` take the ``on_value``, while other locations take the ``off_value``. Parameters ---------- indices: tensor<[D],T> (Required) * Tensor, values indicate the locations for each one-hot vector to take the ``on_value``. one_got_vector_size: i32 (Required) * Indicates the number of returning vectors. axis: const i32 (Optional) * Indicates which dimension to append the new axis. * If the input indices is rank ``D``, the output tensor will have rank ``D+1``. * Default to ``-1`` (the last dimension). on_value: const i32 (Optional) * Values for locations where defined in ``indices``. * Default to ``1``. off_value: const i32 (Optional) * Default to ``0``. Returns ------- tensor<\*?,T> * A tensor that contains one-hot vectors. Attributes ---------- T: fp32 """ input_spec = InputSpec( indices=IntTensorInputType(), one_hot_vector_size=IntInputType(), axis=IntInputType(const=True, optional=True), on_value=IntOrFloatInputType(const=True, optional=True), off_value=IntOrFloatInputType(const=True, optional=True), ) def default_inputs(self): return DefaultInputs( axis=-1, on_value=1, off_value=0, ) def __init__(self, **kwargs): super(one_hot, self).__init__(**kwargs) def type_inference(self): on_type = self.on_value.dtype off_type = self.off_value.dtype if on_type != off_type: raise TypeError( "Parameters on_value and off_value must have same input types." ) if self.axis.val < -self.indices.rank - 1 or self.axis.val > self.indices.rank: raise IndexError( "Axis value {} is out of bounds for {} node {}".format( self.axis.val, self.op_type, self.name ) ) indices_shape = list(self.indices.shape) depth_value = self.one_hot_vector_size.val if depth_value is None: depth_value = get_new_symbol() elif depth_value < 0: raise ValueError("Parameter one_hot_vector_size must be non-negative") retshape = indices_shape if self.axis.val < 0: cut = len(retshape) + self.axis.val + 1 else: cut = self.axis.val retshape = retshape[0:cut] + [depth_value] + retshape[cut:] return types.tensor(on_type, retshape)
[docs]@register_op(doc_str="") 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)
[docs]@register_op(doc_str="") 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: fp32 """ input_spec = InputSpec( end=IntOrFloatInputType(), start=IntOrFloatInputType(), step=IntOrFloatInputType(), ) def __init__(self, **kwargs): super(range_1d, self).__init__(**kwargs) @precondition(allow=VALUE) def value_inference(self): start = self.start.val end = self.end.val step = self.step.val return np.arange(start, end, step).astype(np.int32) 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)
[docs]@register_op(doc_str="") 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: Any """ input_spec = InputSpec(x=TensorInputType(), reps=TensorInputType(),) def __init__(self, **kwargs): super(tile, self).__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)
[docs]@register_op(doc_str="") class argsort(Operation): """ Returns a tensor containing the indices of the sorted values along a given axis of the input tensor. Parameters ---------- x: <\*?, T> (Required) * Input tensor. * axis: const<i32> (Optional) * Default to ``-1`` (the last dimension). * Axis to perform the operation. * ascending: const<bool> (Optional) * Default to ``False``, sort in descending order. ``True`` to sort in ascending order. Returns ------- tensor<\*?, int32> * Tensor containing the indices of the sorted values Attributes ---------- T: fp32 """ input_spec = InputSpec( x=TensorInputType(), axis=IntInputType(const=True, optional=True), ascending=BoolInputType(const=True, optional=True), ) def default_inputs(self): return DefaultInputs( axis=-1, ascending=False, ) def __init__(self, **kwargs): super(argsort, self).__init__(**kwargs) def type_inference(self): return types.tensor(types.int32, self.x.shape) @precondition(allow=VALUE) def value_inference(self): # The default np argsort mode is ascending, which is opposite to MIL's argsort op. if self.ascending.val: return np.argsort(self.x.val, axis=self.axis.val) return np.argsort(-self.x.val, axis=self.axis.val)
[docs]@register_op(doc_str="") class topk(Operation): """ Returns a tensor containing top or bottom ``k`` values and the corresponding indices of the input tensor along a given axis. Parameters ---------- x: <\*?, T> (Required) * Input tensor. k: const<i32> (Optional) * Default to ``1``. * Number of values/indices to be computed along each axis. * axis: const<i32> (Optional) * Defaults to ``-1`` (last dimension). * Axis to perform the operation. * ascending: const<bool> (Optional) * Default to ``False``, sort in descending order. ``True`` to sort in ascending order. Returns ------- tensor<\*?, T> * Values of top/bottom ``k`` elements. tensor<\*?, int32> * Indices of the top/bottom ``k`` elements along axis. Attributes ---------- T: fp32, int32 """ input_spec = InputSpec( x=TensorInputType(), k=IntInputType(const=True, optional=True), axis=IntInputType(const=True, optional=True), ascending=BoolInputType(const=True, optional=True), ) def default_inputs(self): return DefaultInputs( k=1, axis=-1, ascending=False, ) def __init__(self, **kwargs): super(topk, self).__init__(**kwargs) def type_inference(self): x_type = self.x.dtype x_shape = self.x.shape k = self.k.val axis = self.axis.val if not is_symbolic(x_shape[axis]) and k > x_shape[axis]: msg = "K={} is greater than size of the given axis={}" raise ValueError(msg.format(k, axis)) ret_shape = list(x_shape) ret_shape[axis] = k return types.tensor(x_type, ret_shape), types.tensor(types.int32, ret_shape) @precondition(allow=VALUE) def value_inference(self): indices = np.argsort(self.x.val, axis=self.axis.val) if not self.ascending.val: indices = np.argsort(-self.x.val, axis=self.axis.val) slc = [slice(None)] * self.x.rank slc[self.axis.val] = slice(0, self.k.val) indices = indices[tuple(slc)] values = np.take_along_axis(self.x.val, indices, axis=self.axis.val) return values, indices
[docs]@register_op(doc_str="") class flatten2d(Operation): """ Flattens input tensor into 2d tensor by flattening dimensions before and after the provided axis. Parameters ---------- x: tensor<[*d], T> (Required) * Input tensor. * axis: const<f32> (Optional) * Defaults to ``1``. * Negative axis is supported. Returns ------- tensor<d_prior, d_post, T> * ``d_prior`` is product of dimensions ``x[:axis]`` * ``d_post`` is product of dimensions ``x[axis:]`` Examples -------- 1. ``input_shape = (3, ), axis = -1, output_shape = (1, 3)`` 2. ``input_shape = (3, ), axis = 1, output_shape = (3, 1)`` 3. ``input_shape = (4, 3), axis = -1, output_shape = (4, 3)`` 4. ``input_shape = (2, 3, 2), axis = -1, output_shape = (6, 2)`` 5. ``input_shape = (5, 5, 2), axis = 1, output_shape = (5, 10)`` Attributes ---------- T: fp32 """ input_spec = InputSpec( x=TensorInputType(), axis=IntInputType(const=True, optional=True) ) def default_inputs(self): return DefaultInputs( axis=1, ) def __init__(self, **kwargs): super(flatten2d, self).__init__(**kwargs) def type_inference(self): shape = list(self.x.shape) axis = self.axis.val dim_pre_axis = np.prod(shape[:axis]) dim_post_axis = np.prod(shape[axis:]) new_shape = [dim_pre_axis, dim_post_axis] return types.tensor(self.x.dtype, tuple(new_shape)) @precondition(allow=VALUE | SYMBOL) def value_inference(self): shape = self.x.shape axis = self.axis.val dim_pre_axis = np.prod(shape[:axis]) dim_post_axis = np.prod(shape[axis:]) return self.x.val.reshape(dim_pre_axis, dim_post_axis)
[docs]@register_op(doc_str="") class shape(Operation): """ Returns a 1-dimensional tensor with the shape of the input tensor Parameters ---------- x: tensor<[*?], T> (Required) * Input tensor. Returns ------- tensor<K, i32> * Shape of the input tensor. * ``K = x.rank``. Attributes ---------- T: fp32 """ input_spec = InputSpec(x = ScalarOrTensorInputType()) def __init__(self, **kwargs): super(shape, self).__init__(**kwargs) def type_inference(self): input_rank = self.x.rank return types.tensor(types.int32, tuple([input_rank])) def value_inference(self): if any_symbolic(self.x.shape): # convert elements in shape to int32 res = [x if is_symbolic(x) else np.int32(x) for x in self.x.shape] return np.array(res) else: return np.array(self.x.shape).astype(np.int32)
[docs]@register_op(doc_str="") class concat(Operation): """ Concatenates tensors along a dimension. Parameters ---------- values: Tuple[tensor<[d0, d1, ..., d_axis_i, ..., d_n],T>] (Required) * The number of dimensions of the input tensors must match, and all dimensions except ``axis`` must be equal. * The tensors may be variadic, but the number of tensors must be determined at compile time (i.e. a tuple). axis: const<int32> (Required) * The dimension along which to concatenate. Must be in the range ``[-rank(values[i]), rank(values[i]))`` for all ``i``. interleave: const<bool> (Optional, Default=False) * If true, concatenate the inputs by interleaving them. * If true, all the inputs to this op must have the exact same shape. Examples -------- .. sourcecode:: python in1 : shape (3, 2), value = [[1, 2], [3, 4], [5, 6]] in2 : shape (3, 2), value = [[7, 8], [9, 10], [11, 12]] axis = 0 if interleave = False (default) output : shape (6, 2) output[0:3, :] = in1 output[3:6, :] = in2 value = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]] if interleave = True output : shape (6, 2) output[0::2, :] = in1 output[1::2, :] = in2 value = [[1, 2], [7, 8], [3, 4], [9, 10], [5, 6], [11, 12]] Returns ------- tensor<[d0, d1,...d_axis_out, ..., d_n],T> * Where ``d_axis_out = sum(d_axis_i)``. Attributes ---------- T: fp32, int32 """ input_spec = InputSpec(values=TupleInputType(), axis=IntInputType(const=True), interleave=BoolInputType(const=True, optional=True)) def default_inputs(self): return DefaultInputs( interleave=False, ) def __init__(self, **kwargs): super(concat, self).__init__(**kwargs) def type_inference(self): concat_dim_len = 0 if len(self.values) == 0: raise ValueError("Concat {} got 0 values".format(self.name)) # Validate values have the same rank rank = self.values[0].rank for v in self.values: if v.rank != rank: msg = "Input {} has rank {} != other inputs rank {}" raise ValueError(msg.format(v.name, v.rank, rank)) # Check concat axis is within (-rank, rank) concat_axis = self.axis.val if concat_axis < 0: concat_axis += rank if rank > 0 and (concat_axis < 0 or concat_axis >= rank): msg = "In {} of op_type {}: axis out of bound for input " + "(rank {})" raise ValueError(msg.format(self.name, self.op_type, rank)) # Validate primitive types are compatible dtype = self.values[0].dtype for v in self.values[1:]: new_dtype = promoted_primitive_type(v.dtype, dtype) if new_dtype is None: msg = "Incompatible primitive types concat: {} vs {}" raise ValueError(msg.format(v.dtype, dtype)) dtype = new_dtype # validate that non-axis dimensions match retshape = list(self.values[0].shape) for v in self.values[1:]: for i in range(rank): if is_symbolic(retshape[i]) or is_symbolic(v.shape[i]): continue if i != concat_axis and retshape[i] != v.shape[i]: msg = 'Dimension mismatch in {} ("{}"): shapes {} vs. {}' raise ValueError( msg.format(self.op_type, self.name, retshape, v.shape) ) if self.interleave.val and retshape[i] != v.shape[i]: msg = 'Dimension mismatch in {} ("{}"): shapes {} vs. {}. ' \ 'All inputs must have same shape when \'interleave\' option is True.' raise ValueError( msg.format(self.op_type, self.name, retshape, v.shape) ) # Get length of concat dim concat_dim_len = 0 for v in self.values: if len(v.shape) == 0: taxis = 1 else: taxis = v.shape[concat_axis] if is_symbolic(taxis): concat_dim_len = get_new_symbol() break concat_dim_len += taxis if len(retshape) == 0: retshape = [concat_dim_len] else: retshape[concat_axis] = concat_dim_len return types.tensor(dtype, retshape) @precondition(allow=VALUE | SYMBOL | NONE) def value_inference(self): values = [] for v in self.values: if v.sym_val is not None: values.append(v.sym_val) continue if v.rank == 0: values.append(get_new_symbol()) continue if any_symbolic(v.shape): values.append(None) continue # we support value inference when number of elements for each tensor is less than 10 shape = v.shape num_element = np.prod(shape) if num_element > 10: values.append(None) continue symbolic_tensor = [get_new_symbol() for _ in range(num_element)] symbolic_tensor = np.reshape(np.array(symbolic_tensor), shape) values.append(symbolic_tensor) if any([val is None for val in values]): return None if not isinstance(values[0], np.ndarray) or values[0].shape == (): return np.stack(values, axis=self.axis.val) return np.concatenate(values, axis=self.axis.val)
[docs]@register_op(doc_str="") 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))
[docs]@register_op(doc_str="") class stack(Operation): """ Concatenates tensors along a dimension. Parameters ---------- values: Tuple[tensor<[d0, d1,...d_axis_i, ..., d_n],T>] (Required) * All tensors must have identical shape. axis: const<i32> (Required) * The dimension along which to concatenate. Must be in the range ``[-rank(values[i]), rank(values[i]))`` for all ``i``. Returns ------- tenor<[d0, d1,...d_axis_out, ..., d_n],T> * Where ``d_axis_out = sum(d_axis_i)``. Attributes ---------- T: fp32 """ input_spec = InputSpec(values=TupleInputType(), axis=IntInputType(const=True),) def __init__(self, **kwargs): super(stack, self).__init__(**kwargs) def type_inference(self): num_tensors = len(self.values) if num_tensors == 0: raise ValueError("Cannot stack 0 tensor") # get the first value without symbolic shape t_shape = None for value in self.values: if not any_symbolic(value.shape): t_shape = value.shape break t_shape = self.values[0].shape if t_shape is None else t_shape # compare all shape for t in self.values: if not is_compatible_symbolic_vector(t.shape, t_shape): msg = "Component tensor {} has shape {}, others have {}" raise ValueError(msg.format(t.name, t.shape, t_shape)) ret_shape = list(t_shape) ret_shape.insert(self.axis.val, num_tensors) return types.tensor(self.values[0].dtype, ret_shape) @precondition(allow=VALUE | SYMBOL | NONE) def value_inference(self): is_all_rank_zero = all([v.rank == 0 for v in self.values]) values = [ v.sym_val if v.sym_val is not None else get_new_symbol() for v in self.values ] if any([is_symbolic(v) for v in values]) and not is_all_rank_zero: return None return np.stack(values, self.axis.val)
# identity is used for renaming and is rarely necessary. See # `loop_invariant_elimination` pass for a rare use case.
[docs]@register_op(doc_str="") class identity(Operation): """ Returns a tensor with the same shape and contents as input. Parameters ---------- x: tensor<\*?, T> (Required) * Input tensor. Returns ------- tensor<\*?, T> * Same type and shape as the input tensor. """ input_spec = InputSpec(x=ListOrScalarOrTensorInputType()) def __init__(self, **kwargs): super(identity, self).__init__(**kwargs) def type_inference(self): return self.x.sym_type @precondition(allow=VALUE | SYMBOL) def value_inference(self): return self.x.sym_val