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
