tvm.relay.op

Relay core operators.

The base node types for the Relay language.

class tvm.relay.op.op.Op

A Relay operator definition.

get_attr(attr_name)

Get additional attribute about the operator.

Parameters:attr_name (str) – The attribute name.
Returns:value – The attribute value
Return type:object
class tvm.relay.op.op.OpPattern

Operator generic patterns

See also

top.tag
Contains explanation of the tag type.
tvm.relay.op.op.debug(expr, debug_func=None)

The main entry point to the debugger.

tvm.relay.op.op.get(op_name)

Get the Op for a given name

Parameters:op_name (str) – The operator name
Returns:op – The op of the corresponding name
Return type:Op
tvm.relay.op.op.register(op_name, attr_key, value=None, level=10)

Register an operator property of an operator.

Parameters:
  • op_name (str) – The name of operator
  • attr_key (str) – The attribute name.
  • value (object, optional) – The value to set
  • level (int, optional) – The priority level
Returns:

fregister – Register function if value is not specified.

Return type:

function

tvm.relay.op.op.register_alter_op_layout(op_name, alter_layout=None, level=10)

Register alter op layout function for an op

Parameters:
  • op_name (str) – The name of the operator
  • alter_layout (function (attrs: Attrs, inputs: List[Expr]) -> new_expr: Expr) – The function for changing the layout or replacing the operator
  • level (int) – The priority level
tvm.relay.op.op.register_compute(op_name, compute=None, level=10)

Register compute function for an op.

Parameters:
  • op_name (str) – The name of the op.
  • compute (function (attrs: Attrs, inputs: List[Tensor], out_type: Type, target:Target)) – -> List[Tensor] The compute function.
  • level (int) – The priority level
tvm.relay.op.op.register_gradient(op_name, fgradient, level=10)

Register operator pattern for an op.

Parameters:
  • op_name (str) – The name of the op.
  • fgradient (function (orig_expr : Expr, output_grad : Expr) -> new_expr : Expr) – The gradient being used.
  • level (int) – The priority level
tvm.relay.op.op.register_pattern(op_name, pattern, level=10)

Register operator pattern for an op.

Parameters:
  • op_name (str) – The name of the op.
  • pattern (int) – The pattern being used.
  • level (int) – The priority level
tvm.relay.op.op.register_schedule(op_name, schedule=None, level=10)

Register schedule function for an op

Parameters:
  • op_name (str) – The name of the op.
  • schedule (function (attrs: Attrs, outs: List[Tensor], target: Target) -> sch: Schedule) – The schedule function.
  • level (int) – The priority level
tvm.relay.op.op.schedule_injective(attrs, outputs, target)

Generic schedule for binary broadcast.

Reduce operators.

tvm.relay.op.reduce.argmax(data, axis=None, keepdims=False, exclude=False)

Returns the indices of the maximum values along an axis.

Parameters:
  • data (relay.Expr) – The input data
  • axis (None or int or tuple of int) – Axis or axes along which a argmax operation is performed. The default, axis=None, will find the indices of the maximum element of the elements of the input array. If axis is negative it counts from the last to the first axis.
  • keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input array.
  • exclude (bool) – If exclude is true, reduction will be performed on the axes that are NOT in axis instead.
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.reduce.argmin(data, axis=None, keepdims=False, exclude=False)

Returns the indices of the minimum values along an axis.

Parameters:
  • data (relay.Expr) – The input data
  • axis (None or int or tuple of int) – Axis or axes along which a argmin operation is performed. The default, axis=None, will find the indices of minimum element all of the elements of the input array. If axis is negative it counts from the last to the first axis.
  • keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input array.
  • exclude (bool) – If exclude is true, reduction will be performed on the axes that are NOT in axis instead.
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.reduce.max(data, axis=None, keepdims=False, exclude=False)

Computes the max of array elements over given axes.

Parameters:
  • data (relay.Expr) – The input data
  • axis (None or int or tuple of int) – Axis or axes along which the max operation is performed. The default, axis=None, will find the max element from all of the elements of the input array. If axis is negative it counts from the last to the first axis.
  • keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input array.
  • exclude (bool) – If exclude is true, reduction will be performed on the axes that are NOT in axis instead.
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.reduce.mean(data, axis=None, keepdims=False, exclude=False)

Computes the mean of array elements over given axes.

Parameters:
  • data (relay.Expr) – The input data
  • axis (None or int or tuple of int) – Axis or axes along which a mean operation is performed. The default, axis=None, will find the indices of minimum element all of the elements of the input array. If axis is negative it counts from the last to the first axis.
  • keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input array.
  • exclude (bool) – If exclude is true, reduction will be performed on the axes that are NOT in axis instead.
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.reduce.min(data, axis=None, keepdims=False, exclude=False)

Computes the min of array elements over given axes.

Parameters:
  • data (relay.Expr) – The input data
  • axis (None or int or tuple of int) – Axis or axes along which a minimum operation is performed. The default, axis=None, will find the minimum element from all of the elements of the input array. If axis is negative it counts from the last to the first axis.
  • keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input array.
  • exclude (bool) – If exclude is true, reduction will be performed on the axes that are NOT in axis instead.
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.reduce.prod(data, axis=None, keepdims=False, exclude=False)

Computes the products of array elements over given axes.

Parameters:
  • data (relay.Expr) – The input data
  • axis (None or int or tuple of int) – Axis or axes along which a product is performed. The default, axis=None, will find the indices of minimum element all of the elements of the input array. If axis is negative it counts from the last to the first axis.
  • keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input array.
  • exclude (bool) – If exclude is true, reduction will be performed on the axes that are NOT in axis instead.
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.reduce.sum(data, axis=None, keepdims=False, exclude=False)

Computes the sum of array elements over given axes.

Parameters:
  • data (relay.Expr) – The input data
  • axis (None or int or tuple of int) – Axis or axes along which a sum is performed. The default, axis=None, will sum all of the elements of the input array. If axis is negative it counts from the last to the first axis.
  • keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input array.
  • exclude (bool) – If exclude is true, reduction will be performed on the axes that are NOT in axis instead.
Returns:

result – The computed result.

Return type:

relay.Expr

Basic tensor operations.

tvm.relay.op.tensor.abs(data)

Compute element-wise absolute of data.

Parameters:data (relay.Expr) – The input data
Returns:result – The computed result.
Return type:relay.Expr
tvm.relay.op.tensor.add(lhs, rhs)

Addition with numpy-style broadcasting.

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

Examples

x = relay.Var("a") # shape is [2, 3]
y = relay.Var("b") # shape is [2, 1]
z = relay.add(x, y)  # result shape is [2, 3]
tvm.relay.op.tensor.ceil(data)

Compute element-wise ceil of data.

Parameters:data (relay.Expr) – The input data
Returns:result – The computed result.
Return type:relay.Expr
tvm.relay.op.tensor.clip(a, a_min, a_max)

Clip the elements in a between a_min and a_max. a_min and a_max are cast to a’s dtype.

Parameters:
  • a (relay.Expr) – The input tensor.
  • a_min (float) – The clip minimum.
  • a_max (float) – The clip maximum.
Returns:

resulta with elements clipped between a_min and a_max.

Return type:

relay.Expr

Examples

tvm.relay.op.tensor.concatenate(data, axis)

Concatenate the input tensors along the given axis.

Parameters:
  • data (Union(List[relay.Expr], Tuple[relay.Expr])) – A list of tensors.
  • axis (int) – The axis along which the tensors are concatenated.
Returns:

result – The concatenated tensor.

Return type:

relay.Expr

tvm.relay.op.tensor.copy(data)

Copy a tensor.

Parameters:data (relay.Expr) – The tensor to be copied.
Returns:result – The copied result.
Return type:relay.Expr
tvm.relay.op.tensor.device_copy(data, src_dev, dst_dev)

Copy data from the source device to the destination device. This operator helps data transferring between difference contexts for heterogeneous execution.

Parameters:
  • data (tvm.relay.Expr) – The tensor to be copied.
  • src_dev (Union[TVMContext, str]) – The source device where the data is copied from.
  • dst_dev (Union[TVMContext, str]) – The destination device where the data is copied to.
Returns:

result – The copied result.

Return type:

tvm.relay.Expr

tvm.relay.op.tensor.divide(lhs, rhs)

Division with numpy-style broadcasting.

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.equal(lhs, rhs)

Broadcasted elementwise test for (lhs == rhs).

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.exp(data)

Compute elementwise exp of data.

Parameters:data (relay.Expr) – The input data
Returns:result – The computed result.
Return type:relay.Expr
tvm.relay.op.tensor.floor(data)

Compute element-wise floor of data.

Parameters:data (relay.Expr) – The input data
Returns:result – The computed result.
Return type:relay.Expr
tvm.relay.op.tensor.greater(lhs, rhs)

Broadcasted elementwise test for (lhs > rhs).

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.greater_equal(lhs, rhs)

Broadcasted elementwise test for (lhs >= rhs).

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.left_shift(lhs, rhs)

Left shift with numpy-style broadcasting.

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.less(lhs, rhs)

Broadcasted elementwise test for (lhs < rhs).

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.less_equal(lhs, rhs)

Broadcasted elementwise test for (lhs <= rhs).

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.log(data)

Compute elementwise log of data.

Parameters:data (relay.Expr) – The input data
Returns:result – The computed result.
Return type:relay.Expr
tvm.relay.op.tensor.maximum(lhs, rhs)

Maximum with numpy-style broadcasting.

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.minimum(lhs, rhs)

Minimum with numpy-style broadcasting.

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.mod(lhs, rhs)

Mod with numpy-style broadcasting.

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.multiply(lhs, rhs)

Multiplication with numpy-style broadcasting.

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.negative(data)

Compute element-wise negative of data.

Parameters:data (relay.Expr) – The input data
Returns:result – The computed result.
Return type:relay.Expr
tvm.relay.op.tensor.not_equal(lhs, rhs)

Broadcasted elementwise test for (lhs != rhs).

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.ones(shape, dtype)

Fill array with ones.

Parameters:
  • shape (tuple of int) – The shape of the target.
  • dtype (data type) – The data type of the target.
Returns:

result – The resulting tensor.

Return type:

relay.Expr

tvm.relay.op.tensor.ones_like(data)

Returns an array of ones, with same type and shape as the input.

Parameters:data (relay.Expr) – The input data
Returns:result – The computed result.
Return type:relay.Expr
tvm.relay.op.tensor.power(lhs, rhs)

Power with numpy-style broadcasting.

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.right_shift(lhs, rhs)

Right shift with numpy-style broadcasting.

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.round(data)

Compute element-wise round of data.

Parameters:data (relay.Expr) – The input data
Returns:result – The computed result.
Return type:relay.Expr
tvm.relay.op.tensor.sigmoid(data)

Compute elementwise sigmoid of data.

Parameters:data (relay.Expr) – The input data
Returns:result – The computed result.
Return type:relay.Expr
tvm.relay.op.tensor.sqrt(data)

Compute elementwise sqrt of data.

Parameters:data (relay.Expr) – The input data
Returns:result – The computed result.
Return type:relay.Expr
tvm.relay.op.tensor.subtract(lhs, rhs)

Subtraction with numpy-style broadcasting.

Parameters:
  • lhs (relay.Expr) – The left hand side input data
  • rhs (relay.Expr) – The right hand side input data
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.tensor.tanh(data)

Compute element-wise tanh of data.

Parameters:data (relay.Expr) – The input data
Returns:result – The computed result.
Return type:relay.Expr
tvm.relay.op.tensor.trunc(data)

Compute element-wise trunc of data.

Parameters:data (relay.Expr) – The input data
Returns:result – The computed result.
Return type:relay.Expr
tvm.relay.op.tensor.zeros(shape, dtype)

Fill array with zeros.

Parameters:
  • shape (tuple of int) – The shape of the target.
  • dtype (data type) – The data type of the target.
Returns:

result – The resulting tensor.

Return type:

relay.Expr

tvm.relay.op.tensor.zeros_like(data)

Returns an array of zeros, with same type and shape as the input.

Parameters:data (relay.Expr) – The input data
Returns:result – The computed result.
Return type:relay.Expr

Transform operators.

tvm.relay.op.transform.broadcast_to(data, shape)

Return an scalar value array with the same type, broadcast to the provided shape.

Parameters:
  • data (relay.Expr) – The input tensor.
  • shape (shape) – Provide the shape to broadcast to.
Returns:

result – The resulting tensor.

Return type:

relay.Expr

tvm.relay.op.transform.broadcast_to_like(data, broadcast_type)

Return an scalar value array with the same shape and type as the input array.

Parameters:
  • data (relay.Expr) – The input tensor.
  • broadcast_type (relay.Expr) – Provide the type to broadcast to.
Returns:

result – The resulting tensor.

Return type:

relay.Expr

tvm.relay.op.transform.cast(data, dtype)

Cast input tensor to data type.

Parameters:
  • data (relay.Expr) – The input data to the operator.
  • dtype (str) – The target data type
Returns:

result – The casted result.

Return type:

relay.Expr

tvm.relay.op.transform.collapse_sum_like(data, collapse_type)

Return an scalar value array with the same shape and type as the input array.

Parameters:
  • data (relay.Expr) – The input tensor.
  • collapse_type (relay.Expr) – Provide the type to collapse to.
Returns:

result – The resulting tensor.

Return type:

relay.Expr

tvm.relay.op.transform.expand_dims(data, axis, num_newaxis=1)

Insert num_newaxis axises at the position given by axis.

Parameters:
  • data (relay.Expr) – The input data to the operator.
  • axis (int) – The axis at which the input array is expanded. Should lie in range [-data.ndim - 1, data.ndim]. If axis < 0, it is the first axis inserted; If axis >= 0, it is the last axis inserted in Python’s negative indexing.
  • num_newaxis (int) – Number of axes to be inserted. Should be >= 0.
Returns:

result – The reshaped result.

Return type:

relay.Expr

tvm.relay.op.transform.full(fill_value, shape=(), dtype='')

Fill array with scalar value.

Parameters:
  • fill_value (relay.Expr) – The value to fill. Must be a scalar.
  • shape (tuple of int) – The shape of the target.
  • dtype (data type, optional (defaults to data type of the fill value)) – The data type of the target.
Returns:

result – The resulting tensor.

Return type:

relay.Expr

tvm.relay.op.transform.full_like(data, fill_value)

Return an scalar value array with the same shape and type as the input array.

Parameters:
  • data (relay.Expr) – The input tensor.
  • fill_value (relay.Expr) – The scalar value to fill.
Returns:

result – The resulting tensor.

Return type:

relay.Expr

tvm.relay.op.transform.layout_transform(data, src_layout, dst_layout)

Transform the layout of a tensor

Parameters:
  • data (relay.Expr) – The source tensor to be transformed
  • src_layout (str) – The source layout. (e.g NCHW)
  • dst_layout (str) – The destination layout. (e.g. NCHW16c)
Returns:

ret – The transformed tensor.

Return type:

relay.Expr

tvm.relay.op.transform.reshape(data, newshape)

Reshapes the input array.

Example:

To give user more convenience in without doing manual shape inference, some dimensions of the shape can take special values from the set {0, -1, -2, -3, -4}. The significance of each is explained below:

  • 0 copy this dimension from the input to the output shape.

Example:

- data.shape = (2,3,4), newshape = (4,0,2), result.shape = (4,3,2)
- data.shape = (2,3,4), newshape = (2,0,0), result.shape = (2,3,4)
  • -1 infers the dimension of the output shape by using the remainder of the input dimensions

keeping the size of the new array same as that of the input array. At most one dimension of shape can be -1.

Example:

- data.shape = (2,3,4), newshape = (6,1,-1), result.shape = (6,1,4)
- data.shape = (2,3,4), newshape = (3,-1,8), result.shape = (3,1,8)
- data.shape = (2,3,4), newshape = (-1,), result.shape = (24,)
  • -2 copy all/remainder of the input dimensions to the output shape.

Example:

- data.shape = (2,3,4), newshape = (-2,), result.shape = (2,3,4)
- data.shape = (2,3,4), newshape = (2,-2), result.shape = (2,3,4)
- data.shape = (2,3,4), newshape = (-2,1,1), result.shape = (2,3,4,1,1)
  • -3 use the product of two consecutive dimensions of the input shape

as the output dimension.

Example:

- data.shape = (2,3,4), newshape = (-3,4), result.shape = (6,4)
- data.shape = (2,3,4,5), newshape = (-3,-3), result.shape = (6,20)
- data.shape = (2,3,4), newshape = (0,-3), result.shape = (2,12)
- data.shape = (2,3,4), newshape = (-3,-2), result.shape = (6,4)
  • -4 split one dimension of the input into two dimensions passed subsequent

to -4 in shape (can contain -1).

Example:

- data.shape = (2,3,4), newshape = (-4,1,2,-2), result.shape =(1,2,3,4)
- data.shape = (2,3,4), newshape = (2,-4,-1,3,-2), result.shape = (2,1,3,4)
Parameters:
  • data (relay.Expr) – The input data to the operator.
  • newshape (Union[int, Tuple[int], List[int]]) – The new shape. Should be compatible with the original shape.
Returns:

result – The reshaped result.

Return type:

relay.Expr

tvm.relay.op.transform.reshape_like(data, shape_like)

Reshapes the input array by the size of another array. For an input array with shape (d1, d2, ..., dk), reshape_like operation reshapes the input array into an output array with the same shape as the second input array. .. note:: Sizes for both array should be compatible.

Parameters:
  • data (relay.Expr) – The input data to the operator.
  • shape_like (tuple of int) – The new shape. Should be compatible with the original shape.
Returns:

ret – The computed result.

Return type:

relay.Expr

tvm.relay.op.transform.slice_like(data, shape_like, axes=None)

Slice the first input with respect to the second input.

For an input array with shape (d1, d2, ..., dk), slice_like operation slices the the input array corresponding size of second array. By default will slice on all axes.

Parameters:
  • data (tvm.relay.Expr) – The source array.
  • shape_like (tvm.relay.Expr) – The new shape.
  • axes (Optional[Tuple[int]]) – List of axes on which input data will be sliced according to the corresponding size of the second input. By default will slice on all axes. Negative axes mean counting in reverse.
Returns:

result – The computed result.

Return type:

relay.Expr

tvm.relay.op.transform.split(data, indices_or_sections, axis=0)

Split input tensor along axis by sections or indices.

If indices_or_sections is an integer, the input will be divided equally along given axis. If such a split is not possible, an error is raised.

If indices_or_sections is a tuple of sorted integers, the entries indicate where along axis the array is split.

Parameters:
  • data (relay.Expr) – The source array.
  • indices_or_sections (int or tuple of int) – Indices or sections to split into. Accepts an int or a tuple
  • axis (int, optional) – The axis over which to split.
Returns:

ret – The computed result.

Return type:

relay.Tuple([relay.Expr, relay.Expr])

tvm.relay.op.transform.squeeze(data, axis=None)

Squeeze axes in the array.

Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • axis (None or List[int]) – The set of axes to remove. If axis = None, remove all axis of dimensions 1. If any specified axis has dimension that does not equal 1, it is an error.
Returns:

result – The squeezed result.

Return type:

tvm.relay.Expr

tvm.relay.op.transform.strided_slice(data, begin, end, strides=None)

Strided slice of an array..

Parameters:
  • data (relay.Expr) – The source array to be sliced.
  • begin (list of int) – The indices to begin with in the slicing.
  • end (list of int) – Indicies indicating end of the slice.
  • strides (list of int, optional) – Specifies the stride values, it can be negative in that case, the input tensor will be reversed in that particular axis.
Returns:

ret – The computed result.

Return type:

relay.Expr

tvm.relay.op.transform.take(data, indices, axis=None)

Take elements from an array along an axis.

Parameters:
  • data (relay.Expr) – The source array.
  • indices (rely.Expr) – The indices of the values to extract.
  • axis (int, optional) – The axis over which to select values. By default, the flattened input array is used.
Returns:

ret – The computed result.

Return type:

relay.Expr

tvm.relay.op.transform.transpose(data, axes=None)

Permutes the dimensions of an array.

Parameters:
  • data (relay.Expr) – The input data to the operator.
  • axes (None or List[int]) – The target axes order, reverse order if not specified.
Returns:

result – The transposed result.

Return type:

relay.Expr

tvm.relay.op.transform.where(condition, x, y)

Selecting elements from either x or y depending on the value of the condition.

Parameters:
  • condition (relay.Expr) – The condition array. The n-th element in y is selected when the n-th value in the condition array is zero. Otherwise, the corresponding element from x will be picked.
  • x (relay.Expr) – The first array to be selected.
  • y (relay.Expr) – The second array to be selected.
Returns:

result – The selected array.

Return type:

relay.Expr

Examples

x = [[1, 2], [3, 4]]
y = [[5, 6], [7, 8]]
condition = [[0, 1], [-1, 0]]
relay.where(conditon, x, y) = [[5, 2], [3, 8]]

condition = [1, 0]
relay.where(conditon, x, y) = [[1, 2], [7, 8]]

Note that the shape of condition, x, and y needs to be the same.

Neural network operations.

tvm.relay.op.nn.nn.avg_pool2d(data, pool_size=(1, 1), strides=(1, 1), padding=(0, 0), layout='NCHW', ceil_mode=False, count_include_pad=False)

2D average pooling operator.

This operator takes data as input and does 2D average value calculation with in pool_size sized window by striding defined by stride

In the default case, where the data_layout is NCHW a data Tensor with shape (batch_size, in_channels, height, width), to produce an output Tensor with the following rule:

with data of shape (b, c, h, w), pool_size (kh, kw)

\[\mbox{out}(b, c, y, x) = \frac{1}{kh * kw} \sum_{m=0}^{kh-1} \sum_{n=0}^{kw-1} \mbox{data}(b, c, \mbox{stride}[0] * y + m, \mbox{stride}[1] * x + n)\]

Padding is applied to data before the computation. ceil_mode is used to take ceil or floor while computing out shape. count_include_pad indicates including or excluding padded input values in computation. This operator accepts data layout specification.

Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • strides (tuple of int, optional) – The strides of pooling.
  • padding (tuple of int, optional) – The padding for pooling.
  • layout (str, optional) – Layout of the input.
  • ceil_mode (bool, optional) – To enable or disable ceil while pooling.
  • count_include_pad (bool, optional) – To include padding to compute the average.
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.batch_flatten(data)

BatchFlatten.

This operator flattens all the dimensions except for the batch dimension. which results a 2D output.

For data with shape (d1, d2, ..., dk) batch_flatten(data) returns reshaped output of shape (d1, d2*...*dk).

Parameters:data (tvm.relay.Expr) – The input data to the operator.
Returns:result – The Flattened result.
Return type:tvm.relay.Expr
tvm.relay.op.nn.nn.batch_norm(data, gamma, beta, moving_mean, moving_var, axis=1, epsilon=1e-05, center=True, scale=True)

Batch normalization layer (Ioffe and Szegedy, 2014). Normalizes the input at each batch, i.e. applies a transformation that maintains the mean activation close to 0 and the activation standard deviation close to 1.

\[\begin{split}data\_mean[i] = mean(data[:,i,:,...]) \\ data\_var[i] = var(data[:,i,:,...])\end{split}\]

Then compute the normalized output, which has the same shape as input, as following:

\[out[:,i,:,...] = \frac{data[:,i,:,...] - data\_mean[i]}{\sqrt{data\_var[i]+\epsilon}} * gamma[i] + beta[i]\]

Both mean and var returns a scalar by treating the input as a vector.

Assume the input has size k on axis 1, then both gamma and beta have shape (k,).

Besides the inputs and the outputs, this operator accepts two auxiliary states, moving_mean and moving_var, which are k-length vectors. They are global statistics for the whole dataset, which are updated by:

moving_mean = moving_mean * momentum + data_mean * (1 - momentum) moving_var = moving_var * momentum + data_var * (1 - momentum)

The parameter axis specifies which axis of the input shape denotes the ‘channel’ (separately normalized groups). The default is 1. Specifying -1 sets the channel axis to be the last item in the input shape.

Note

This operator can be optimized away for inference.

Parameters:
  • data (tvm.relay.Expr) – Input to which batch_norm will be applied.
  • gamma (tvm.relay.Expr) – The gamma scale factor.
  • beta (tvm.relay.Expr) – The beta offset factor.
  • moving_mean (tvm.relay.Expr) – Running mean of input,
  • moving_var (tvm.relay.Expr) – Running variance of input.
  • axis (int, optional, default=1) – Specify along which shape axis the channel is specified.
  • epsilon (double, optional, default=1e-5) – Small float added to variance to avoid diving by zero.
  • center (boolean, optional, default=True) – If True, add offset of beta to normalized tensor, If False, beta is ignored.
  • scale (boolean, optional, default=True) – If true, multiply by gamma. If False, gamma is not used. When the next layer is piecewise linear (also e.g. nn.relu), this can be disabled since the scaling will be done by the next layer.
Returns:

result – Tuple of normed data (same shape as input), new running mean (k-length vector), and new running variance (k-length vector)

Return type:

relay.Tuple([tvm.relay.Expr, tvm.relay.Expr, tvm.relay.Expr])

tvm.relay.op.nn.nn.bias_add(data, bias, axis=1)

add_bias operator.

Add 1D bias to the axis of data. This function is a special case of add which allows inference of shape of the bias from data.

Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • bias (tvm.relay.Expr) – The bias to be added.
  • axis (int, optional) – The axis to add the bias.
Returns:

result – The final result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.contrib_conv2d_winograd_weight_transform(weight, tile_size)

Weight Transformation part for 2D convolution with winograd algorithm.

We separate this as a single op to enable pre-compute for inference. Use this together with nn.contrib_conv2d_winograd_without_weight_transform

Parameters:
  • weight (tvm.relay.Expr) – The weight expressions.
  • tile_size (int) – The Tile size of winograd. E.g. 2 for F(2x2, 3x3) and 4 for F(4x4, 3x3)
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.contrib_conv2d_winograd_without_weight_transform(data, weight, tile_size, strides=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1, channels=None, kernel_size=None, data_layout='NCHW', kernel_layout='OIHW', out_layout='', out_dtype='')

2D convolution with winograd algorithm.

The basic parameters are the same as the ones in vanilla conv2d. It assumes the weight is pre-transformed by nn.contrib_conv2d_winograd_weight_transform

Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • weight (tvm.relay.Expr) – The weight expressions.
  • tile_size (int) – The Tile size of winograd. E.g. 2 for F(2x2, 3x3) and 4 for F(4x4, 3x3)
  • strides (tuple of int, optional) – The strides of convoltution.
  • padding (tuple of int, optional) – The padding of convolution on both sides of inputs before convolution.
  • dilation (tuple of int, optional) – Specifies the dilation rate to be used for dilated convolution.
  • groups (int, optional) – Number of groups for grouped convolution.
  • channels (int, optional) – Number of output channels of this convolution.
  • kernel_size (tuple of int, optional) – The spatial of the convolution kernel.
  • data_layout (str, optional) – Layout of the input.
  • kernel_layout (str, optional) – Layout of the weight.
  • out_layout (str, optional) – Layout of the output, by default, out_layout is the same as data_layout
  • out_dtype (str, optional) – Specifies the output data type for mixed precision conv2d.
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.conv2d(data, weight, strides=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1, channels=None, kernel_size=None, data_layout='NCHW', kernel_layout='OIHW', out_layout='', out_dtype='')

2D convolution.

This operator takes the weight as the convolution kernel and convolves it with data to produce an output.

In the default case, where the data_layout is NCHW and kernel_layout is OIHW, conv2d takes in a data Tensor with shape (batch_size, in_channels, height, width), and a weight Tensor with shape (channels, in_channels, kernel_size[0], kernel_size[1]) to produce an output Tensor with the following rule:

\[\mbox{out}[b, c, y, x] = \sum_{dy, dx, k} \mbox{data}[b, k, \mbox{strides}[0] * y + dy, \mbox{strides}[1] * x + dx] * \mbox{weight}[c, k, dy, dx]\]

Padding and dilation are applied to data and weight respectively before the computation. This operator accepts data layout specification. Semantically, the operator will convert the layout to the canonical layout (NCHW for data and OIHW for weight), perform the computation, then convert to the out_layout.

Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • weight (tvm.relay.Expr) – The weight expressions.
  • strides (tuple of int, optional) – The strides of convoltution.
  • padding (tuple of int, optional) – The padding of convolution on both sides of inputs before convolution.
  • dilation (tuple of int, optional) – Specifies the dilation rate to be used for dilated convolution.
  • groups (int, optional) – Number of groups for grouped convolution.
  • channels (int, optional) – Number of output channels of this convolution.
  • kernel_size (tuple of int, optional) – The spatial of the convolution kernel.
  • data_layout (str, optional) – Layout of the input.
  • kernel_layout (str, optional) – Layout of the weight.
  • out_layout (str, optional) – Layout of the output, by default, out_layout is the same as data_layout
  • out_dtype (str, optional) – Specifies the output data type for mixed precision conv2d.
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.conv2d_transpose(data, weight, strides=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1, channels=None, kernel_size=None, data_layout='NCHW', kernel_layout='OIHW', output_padding=(0, 0), out_dtype='')

Two dimensional trnasposed convolution operator.

Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • weight (tvm.relay.Expr) – The weight expressions.
  • strides (Tuple[int], optional) – The strides of convoltution.
  • padding (Tuple[int], optional) – The padding of convolution on both sides of inputs.
  • dilation (Tuple[int], optional) – Specifies the dilation rate to be used for dilated convolution.
  • groups (int, optional) – Number of groups for grouped convolution.
  • data_layout (str, optional) – Layout of the input.
  • kernel_layout (str, optional) – Layout of the weight.
  • output_padding (Tuple[int], optional) – Additional zero-padding to be added to one side of the output.
  • out_dtype (str, optional) – Specifies the output data type for mixed precision conv2d.
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.dense(data, weight, units=None)

Dense operator. Applies a linear transformation

\[\]

Y = X * W

Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • weight (tvm.relay.Expr) – The weight expressions.
  • units (int, optional) – Number of hidden units of the dense transformation.
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.dropout(data, rate=0.5)

Applies the dropout operation to the input array.

During training, each element of the input is set to zero with probability p. The whole array is rescaled by 1/(1-p) to keep the expected sum of the input unchanged.

Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • rate (float, optional (default=0.5)) – The probability for an element to be reset to 0.
Returns:

result – The result of dropout

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.global_avg_pool2d(data, layout='NCHW')

2D global average pooling operator.

This operator takes data as input and does 2D average value calculation across each window represented by WxH.

In the default case, where the data_layout is NCHW a data Tensor with shape (batch_size, in_channels, height, width), to produce an output Tensor with the following rule:

with data of shape (b, c, h, w)

\[\mbox{out}(b, c, 1, 1) = \frac{1}{h * w} \sum_{m=0}^{h-1} \sum_{n=0}^{w-1} \mbox{data}(b, c, m, n)\]
Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • layout (str, optional) – Layout of the input.
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.global_max_pool2d(data, layout='NCHW')

2D global maximum pooling operator.

This operator takes data as input and does 2D max value calculation across each window represented by WxH.

In the default case, where the data_layout is NCHW a data Tensor with shape (batch_size, in_channels, height, width), to produce an output Tensor with the following rule:

with data of shape (b, c, h, w)

\[\mbox{out}(b, c, 1, 1) = \max_{m=0, \ldots, h} \max_{n=0, \ldots, w} \mbox{data}(b, c, m, n)\]
Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • layout (str, optional) – Layout of the input.
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.l2_normalize(data, eps, axis=None)

Perform L2 normalization on the input data

\[y(i, j) = x(i, j) / sqrt(max(sum(x^2), eps))\]
Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • eps (float) – epsilon value
  • axis (list of int, optional) – axis over the normalization applied
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.leaky_relu(data, alpha)

This operator takes data as input and does Leaky version of a Rectified Linear Unit.

\[`y = x > 0 ? x : alpha * x`\]
Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • alpha (float) – Slope coefficient for the negative half axis.
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.log_softmax(data, axis=-1)

Computes log softmax.

\[\text{log_softmax}(x)_i = \log \frac{exp(x_i)}{\sum_j exp(x_j)}\]

Note

This operator can be optimized away for inference.

Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • axis (int) – The axis to sum over when computing softmax
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.lrn(data, size=5, axis=1, bias=2, alpha=1e-05, beta=0.75)

This operator takes data as input and does local response normalization.

Normalize the input in a local region across or within feature maps. Each input value is divided by (data / (bias + (alpha * sum_data ^2 /size))^beta) where n is the size of each local region, and the sum is taken over the region centered at that value (zero padding is added where necessary).

\[(data / (bias + (alpha * sum_data ^2 /size))^beta)\]
Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • size (int, optional) – The size of the local region to be considered for normalization.
  • axis (int, optional) – Input data layout channel axis. Default value is 1 for NCHW format
  • bias (float, optional) – The offset parameter to avoid dividing by 0.
  • alpha (float, optional) – The scaling parameter.
  • beta (float, optional) – The exponent parameter.
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.max_pool2d(data, pool_size=(1, 1), strides=(1, 1), padding=(0, 0), layout='NCHW', ceil_mode=False)

2D maximum pooling operator.

This operator takes data as input and does 2D max value calculation with in pool_size sized window by striding defined by stride

In the default case, where the data_layout is NCHW a data Tensor with shape (batch_size, in_channels, height, width), to produce an output Tensor with the following rule:

with data of shape (b, c, h, w) and pool_size (kh, kw)

\[\mbox{out}(b, c, y, x) = \max_{m=0, \ldots, kh-1} \max_{n=0, \ldots, kw-1} \mbox{data}(b, c, \mbox{stride}[0] * y + m, \mbox{stride}[1] * x + n)\]

Padding is applied to data before the computation. ceil_mode is used to take ceil or floor while computing out shape. This operator accepts data layout specification.

Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • strides (tuple of int, optional) – The strides of pooling.
  • padding (tuple of int, optional) – The padding for pooling.
  • layout (str, optional) – Layout of the input.
  • ceil_mode (bool, optional) – To enable or disable ceil while pooling.
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.pad(data, pad_width, pad_value=0.0)

Padding

This operator takes in a tensor and pads each axis by the specified widths using the specified value.

Parameters:
  • data (tvm.relay.Expr) – The input data to the operator
  • pad_width (tuple of <tuple of <int>>, required) – Number of values padded to the edges of each axis, in the format of ((before_1, after_1), …, (before_N, after_N))
  • pad_value (float, optional, default=0.0) – The value used for padding
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.prelu(data, alpha, axis=1)

This operator takes data as input and does Leaky version of a Rectified Linear Unit.

\[`y = x > 0 ? x : alpha * x`\]
Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • alpha (tvm.relay.Expr) – Slope coefficient for the negative half axis.
  • axis (int, optional) – Specify which shape axis the channel is specified.
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.relu(data)

Rectified linear unit.

\[out = max(x, 0)\]
Parameters:data (tvm.relay.Expr) – The input data
Returns:result – The computed result.
Return type:tvm.relay.Expr
tvm.relay.op.nn.nn.softmax(data, axis=-1)

Computes softmax.

\[\text{softmax}(x)_i = \frac{exp(x_i)}{\sum_j exp(x_j)}\]

Note

This operator can be optimized away for inference.

Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • axis (int, optional) – The axis to sum over when computing softmax
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

tvm.relay.op.nn.nn.upsampling(data, scale=1, layout='NCHW', method='NEAREST_NEIGHBOR')

Upsampling.

This operator takes data as input and does 2D scaling to the given scale factor. In the default case, where the data_layout is NCHW with data of shape (n, c, h, w) out will have a shape (n, c, h*scale, w*scale)

method indicates the algorithm to be used while calculating ghe out value and method can be one of (“BILINEAR”, “NEAREST_NEIGHBOR”)

Parameters:
  • data (tvm.relay.Expr) – The input data to the operator.
  • scale (tvm.relay.Expr) – The scale factor for upsampling.
  • layout (str, optional) – Layout of the input.
  • method (str, optional) – Scale method to used [NEAREST_NEIGHBOR, BILINEAR].
Returns:

result – The computed result.

Return type:

tvm.relay.Expr

Multibox operations.

tvm.relay.op.vision.multibox.multibox_prior(data, sizes=(1.0, ), ratios=(1.0, ), steps=(-1.0, -1.0), offsets=(0.5, 0.5), clip=False)

Generate prior(anchor) boxes from data, sizes and ratios.

Parameters:
  • data (relay.Expr) – The input data tensor.
  • sizes (tuple of float, optional) – Tuple of sizes for anchor boxes.
  • ratios (tuple of float, optional) – Tuple of ratios for anchor boxes.
  • steps (Tuple of float, optional) – Priorbox step across y and x, -1 for auto calculation.
  • offsets (tuple of int, optional) – Priorbox center offsets, y and x respectively.
  • clip (boolean, optional) – Whether to clip out-of-boundary boxes.
Returns:

out – 3-D tensor with shape [1, h_in * w_in * (num_sizes + num_ratios - 1), 4]

Return type:

relay.Expr

tvm.relay.op.vision.multibox.multibox_transform_loc(cls_prob, loc_pred, anchor, clip=True, threshold=0.01, variances=(0.1, 0.1, 0.2, 0.2))

Location transformation for multibox detection

Parameters:
  • cls_prob (tvm.relay.Expr) – Class probabilities.
  • loc_pred (tvm.relay.Expr) – Location regression predictions.
  • anchor (tvm.relay.Expr) – Prior anchor boxes.
  • clip (boolean, optional) – Whether to clip out-of-boundary boxes.
  • threshold (double, optional) – Threshold to be a positive prediction.
  • variances (Tuple of float, optional) – variances to be decoded from box regression output.
Returns:

ret

Return type:

tuple of tvm.relay.Expr

Non-maximum suppression operations.

tvm.relay.op.vision.nms.nms(data, valid_count, overlap_threshold=0.5, force_suppress=False, topk=-1)

Non-maximum suppression operator for object detection.

Parameters:
  • data (relay.Expr) – 3-D tensor with shape [batch_size, num_anchors, 6]. The last dimension should be in format of [class_id, score, box_left, box_top, box_right, box_bottom].
  • valid_count (relay.Expr) – 1-D tensor for valid number of boxes.
  • overlap_threshold (float, optional) – Non-maximum suppression threshold.
  • force_suppress (bool, optional) – Suppress all detections regardless of class_id.
  • topk (int, optional) – Keep maximum top k detections before nms, -1 for no limit.
Returns:

out – 3-D tensor with shape [batch_size, num_anchors, 6].

Return type:

relay.Expr