tvm.relax.op#
tvm.relax.op#
Relax core operators.
- tvm.relax.op.abs(x)[源代码]#
Compute element-wise absolute value of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.acos(x)[源代码]#
Compute element-wise arc cos of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.acosh(x)[源代码]#
Compute element-wise arc cosh of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.add(x1, x2)[源代码]#
Addition with numpy-style broadcasting.
Parameters#
- x1Expr
The first input tensor.
- x2Expr
The second input tensor.
Returns#
- resultExpr
The computed result.
Examples#
bb = relax.BlockBuilder() a = relax.Var("a", relax.TensorStructInfo(shape=(2, 3), dtype="float32")) b = relax.Var("b", relax.TensorStructInfo(shape=(2, 1), dtype="float32")) c = bb.normalize(relax.op.add(a, b)) # c has TensorStructInfo(shape=(2, 3), dtype="float32")
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.arange(start, end=None, step=1, dtype=None)[源代码]#
Construct a tensor with evenly spaced elements.
Parameters#
- startUnion[PrimExprLike,PrimValue]
The start of the interval.
- endOptional[Union[PrimExprLike,PrimValue]]
The end of the interval. If not given, it will be set to start, and start will be set to 0.
- stepUnion[PrimExprLike,PrimValue]
The step size.
- dtypeOptional[Union[str, DataType]]
The data type of the created tensor.
Returns#
- resultrelax.Expr
The result tensor.
- tvm.relax.op.argmax(x, axis=None, keepdims=False)[源代码]#
Computes the argmax of tensor elements over given axis.
Parameters#
- xrelax.Expr
The input data tensor
- axisOptional[int]
Axis along which an argmax operation is performed. The default, axis=None, will compute the argmax of all elements in the input tensor. Negative indexing is supported.
- keepdimsbool
If this is set to True, the axis being reduced is left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input tensor.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.argmin(x, axis=None, keepdims=False)[源代码]#
Computes the argmin of tensor elements over given axis.
Parameters#
- xrelax.Expr
The input data tensor
- axisOptional[int]
Axis along which an argmin operation is performed. The default, axis=None, will compute the argmin of all elements in the input tensor. Negative indexing is supported.
- keepdimsbool
If this is set to True, the axis being reduced is left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input tensor.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.argsort(data, axis=-1, descending=False, dtype='int32')[源代码]#
Performs sorting along the given axis and returns an array of indices having same shape as an input array that index data in sorted order.
Parameters#
- datarelax.Expr
The input data tensor.
- axisint
Axis long which to sort the input tensor.
- descendingbool
Whether to sort in descending order, the default is False
- dtypestr
The data type of the output indices.
Returns#
- outrelax.Expr
Tensor with same shape as data.
- tvm.relax.op.asin(x)[源代码]#
Compute element-wise arc sin of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.asinh(x)[源代码]#
Compute element-wise arc sinh of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.assert_op(condition, format_args=None, format='')[源代码]#
Create a call to Relax's assert_op operation (assert is reserved in Python, so the name must be distinct).
Parameters#
- condition: Union[Expr, PrimExpr]
The assertion condition.
- format_args: Optional[Union[Expr, List[Expr]]]
Format arguments for the error message if the condition fails.
- format: Union[str, Expr]
The format string or StringImm for the error message.
Returns#
- resultExpr
A Call to the Relax assert operation.
- tvm.relax.op.astype(x, dtype)[源代码]#
Cast input tensor to the given data type.
Parameters#
- xrelax.Expr
The input data to the operator.
- dtype: Union[str, DataType]
The target data type
Returns#
- resultrelax.Expr
The casted result.
- 参数:
x (RelayExpr)
dtype (str | DataType)
- 返回类型:
RelayExpr
- tvm.relax.op.atan(x)[源代码]#
Compute element-wise arc tan of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.atanh(x)[源代码]#
Compute element-wise arc tanh of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.bitwise_and(x1, x2)[源代码]#
Bitwise AND Parameters ---------- x1 : relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.bitwise_not(x)[源代码]#
Compute bitwise NOT of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.bitwise_or(x1, x2)[源代码]#
Bitwise OR Parameters ---------- x1 : relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.bitwise_xor(x1, x2)[源代码]#
Bitwise XOR Parameters ---------- x1 : relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.broadcast_to(x, shape)[源代码]#
Broadcasts a tensor to a specified shape.
Parameters#
- xrelax.Expr
The input data to the operator.
- shapeUnion[Tuple[PrimExprLike], Expr]
The target shape.
Returns#
- resultrelax.Expr
The broadcasted tensor.
- tvm.relax.op.call_builtin_with_ctx(func, args, *, sinfo_args=None)[源代码]#
Call a builtin function func.
Parameters#
- funcExpr
The builtin function to be called.
- argsExpr
The input arguments.
- sinfo_args: Optional[Union[StructInfo, List[StructInfo]]]
The struct info arguments to the call node.
Returns#
- ret: Call
The created call node.
- 参数:
func (str | RelayExpr)
args (RelayExpr)
sinfo_args (StructInfo | List[StructInfo] | None)
- 返回类型:
- tvm.relax.op.call_dps_packed(func, args, out_sinfo)[源代码]#
Call a destination-passing-style packed function and return the output.
Note: The called function is assumed to be _pure_ (other than modifying the designated output arguments). If the function _does_ result in other side effects, then the compiler may end up removing, reordering, or repeating those effects--no guarantees can be made.
Parameters#
- funcUnion[str, Expr]
The destination-passing-style function, can be ExternFunc.
- argsExpr
The input arguments.
- out_sinfoUnion[TensorStructInfo, List[TensorStructInfo]]
The structure info of the call_dps_packed output. It should be a single or a list of TensorStructInfo. Each one denotes the structure info of a returned tensor.
Returns#
- ret: Call
A call node for the call_dps_packed operator.
- 参数:
func (str | RelayExpr)
args (RelayExpr)
out_sinfo (TensorStructInfo | List[TensorStructInfo])
- 返回类型:
- tvm.relax.op.call_inplace_packed(func, *args, inplace_indices, sinfo_args)[源代码]#
Construct a call to a packed function that consumes some of its arguments "in-place" and returns the mutated arguments (aliased), but should be considered to be otherwise pure. The inplace_indices argument indicates which of the outputs are mutated arguments.
The resulting call will have the same semantics as calling the packed function directly.
Note: This should be used for cases when the user knows that calling the packed function with these arguments will in reality not cause any other side effects. If it is used for a call that does result in other side effects, then the compiler may end up removing, reordering, or repeating that call, with no guarantees made about any side effects from the callee.
Warning: This operator as treated as pure by the type system even though it is performing side effects (mutating some arguments). It is therefore incumbent upon the user to ensure that it is being used safely (viz., that mutated arguments are not live after the mutation, that they do not alias values live after the mutation).
Parameters#
- funcUnion[str, ExternFunc]
The name (global symbol) for a PackedFunc or an ExternFunc node.
- args: Expr
The arguments for the PackedFunc.
- inplace_indicesUnion[int, List[int]]
Specify which arguments should be used for in-place computations. If inplace_indices is a single integer, it will be made into a singleton list. Suppose inplace_indices[i] = j, where j >= 0. Then the i`th output will be an alias of `args[j]. If inplace_indices[i] = -1, then the i`th output will be a freshly allocated tensor. At least one member of `inplace_indices must not be -1.
- sinfo_args: Union[StructInfo, List[StructInfo]]
The list of structure info arguments (giving the structural info for the returned value).
Returns#
- resultExpr
A Relax call, corresponding to call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), sinfo_args)
- 参数:
func (str | ExternFunc | GlobalVar)
args (RelayExpr)
sinfo_args (StructInfo | List[StructInfo])
- 返回类型:
RelayExpr
- tvm.relax.op.call_pure_packed(func, *args, sinfo_args)[源代码]#
Construct a call to a packed function that should be treated as pure, even though packed calls are normally not treated as pure.
The resulting call will have the same semantics as calling the packed function directly.
Note: This should be used for cases when the user knows that calling the packed function with these arguments will in reality not cause any side effects. If it is used for a call that does result in side effects, then the compiler may end up removing, reordering, or repeating that call, with no guarantees made about any side effects from the callee.
Parameters#
- funcUnion[str, ExternFunc]
The name (global symbol) for a PackedFunc or an ExternFunc node.
- args: Expr
The arguments for the PackedFunc.
- sinfo_args: Union[StructInfo, List[StructInfo]]
The list of structure info arguments (giving the structural info for the returned value).
Returns#
- resultExpr
A Relax call, corresponding to call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), sinfo_args)
- 参数:
func (str | ExternFunc | GlobalVar)
args (RelayExpr)
sinfo_args (StructInfo | List[StructInfo])
- 返回类型:
RelayExpr
- tvm.relax.op.call_tir(gvar, args, out_sinfo, tir_vars=None)[源代码]#
Call a tir.prim_func and return the output.
Parameters#
- gvarGlobalVar
The GlobalVar referring to a tir PrimFunc.
- argsExpr
The input arguments.
- out_sinfoUnion[TensorStructInfo, List[TensorStructInfo]]
The structure info of the call_tir output. It should be a single or a list of TensorStructInfo. Each one denotes the structure info of a returned tensor.
- tir_varsOptional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]
ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
Returns#
- ret: Call
A call node for the call_tir operator.
- 参数:
gvar (GlobalVar)
args (RelayExpr)
out_sinfo (TensorStructInfo | List[TensorStructInfo])
tir_vars (ShapeExpr | Tuple[PrimExpr] | List[PrimExpr] | None)
- 返回类型:
- tvm.relax.op.call_tir_inplace(gvar, args, inplace_indices, out_sinfo, tir_vars=None)[源代码]#
Call a TIR PrimFunc and return the result, doing the specified computations in-place (based on the inplace_indices argument; outputs will alias the inputs selected by in-place indices).
Warning: This operator is considered pure by the type system but actually mutates the arguments specified by inplace_indices. This operator should not be used directly, but rather should be inserted by passes that have checked whether it is safe to perform operations in-place (i.e., none of the arguments specified as an output is aliased or is live after calling call_tir_inplace).
Direct calls to this operator should be done for testing purposes only.
Parameters#
- gvarGlobalVar
The GlobalVar referring to a TIR PrimFunc.
- argsExpr
The input arguments.
- inplace_indicesUnion[int, List[int]]
Specify which arguments should be used for in-place computations. If inplace_indices is a single integer, it will be made into a singleton list. Suppose inplace_indices[i] = j, where j >= 0. Then the i`th output will be an alias of `args[j]. If inplace_indices[i] = -1, then the i`th output will be a freshly allocated tensor. At least one member of `inplace_indices must not be -1.
- out_sinfoUnion[TensorStructInfo, List[TensorStructInfo]]
The structure info of the call_tir_inplace output. It should be a single TensorStructInfo or a list of TensorStructInfo. Each one denotes the structure info of a returned tensor. If a list of TensorStructInfo is given, the result will be a tuple of TensorStructInfo.
- tir_varsOptional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]
ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
Returns#
- ret: Call
A call node for the call_tir operator.
- 参数:
gvar (GlobalVar)
args (RelayExpr)
out_sinfo (TensorStructInfo | List[TensorStructInfo])
tir_vars (ShapeExpr | Tuple[PrimExpr] | List[PrimExpr] | None)
- 返回类型:
- tvm.relax.op.call_tir_with_grad(gvar, args, out_sinfo, te_grad_name, te_grad_kwargs=None, tir_vars=None)[源代码]#
Call a tir.prim_func and return the output. This intrinsic will bind a te gradient function (refered by te_grad_name) to the call_tir_with_grad node. The te gradient function will be called by the Gradient pass.
Parameters#
- gvarGlobalVar
The GlobalVar referring to a tir PrimFunc.
- argsExpr
The input arguments.
- out_sinfoUnion[TensorStructInfo, List[TensorStructInfo]]
The structure info of the call_tir_with_grad output. It should be a single or a list of TensorStructInfo. Each one denotes the structure info of a returned tensor.
- te_grad_namestr
The registered name of the te gradient function associated with the call_tir_with_grad node. Must be provided as a keyword argument.
- te_grad_kwargsDict[str, Object], optional
The keyword arguments passed to the te gradient function. Optionally provided as a keyword argument. Default: {}.
- tir_varsOptional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]
ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
Returns#
- ret: Call
A call node for the call_tir_with_grad operator.
- 参数:
gvar (GlobalVar)
args (RelayExpr)
out_sinfo (TensorStructInfo | List[TensorStructInfo])
te_grad_name (str)
tir_vars (ShapeExpr | Tuple[PrimExpr] | List[PrimExpr] | None)
- 返回类型:
- tvm.relax.op.ceil(x)[源代码]#
Take ceil of input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.clip(x, min, max)[源代码]#
Clips tensor values to a specified min and max.
Parameters#
- xrelax.Expr
The input data
- minrelax.Expr
The minimum value
- maxrelax.Expr
The maximum value
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x (RelayExpr)
min (RelayExpr)
max (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.collapse_sum_like(data, collapse_target)[源代码]#
Return a summation of data to the shape of collapse_target.
For details, please see relax.op.collapse_sum_to.
Parameters#
- datarelax.Expr
The input tensor.
- collapse_targetrelax.Expr
The tensor whose shape is the shape to collapse to.
Returns#
- resultrelax.Expr
The result tensor after summation.
- 参数:
data (RelayExpr)
collapse_target (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.collapse_sum_to(data, shape)[源代码]#
Return a summation of data to the given shape.
collapse_sum_to is intended as the backward operator of tvm.relax.op.broadcast_to and other broadcast operators in the automatic differentiation process.
We expect that data is the result of broadcasting some tensor of the given shape in some broadcast operation. Thus the given shape and data.shape must follow broadcast rules.
During computation, all axes of data.shape and shape are checked from right to left. For an axis, if it follows these rules, data will be summed over this axis: - the axis exists in data.shape but not in shape, or - the axis exists in data.shape and equals to 1 in shape.
Parameters#
- datarelax.Expr
The input tensor.
- shapeUnion[Tuple[PrimExprLike], relax.Expr]
The shape to collapse to.
Returns#
- resultrelax.Expr
The result tensor of the given shape after summation.
- tvm.relax.op.concat(tensors, axis=0)[源代码]#
Concatenate the input tensors along the given axis.
Parameters#
- tensorsUnion[relax.Expr, List[relax.Expr]]
An Expr in Tuple type, containing the tensors to be concatenated, or a list of Tensors.
- axisOptional[int]
The axis along which the tensors are concatenated. If axis is None, the input tensor is required to be flattened before concatenation.
Returns#
- result: relax.Expr
The concatenated tensor.
- tvm.relax.op.cos(x)[源代码]#
Compute element-wise cos of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.cosh(x)[源代码]#
Compute element-wise cosh of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.cumprod(data, axis=None, dtype=None, exclusive=False)[源代码]#
Numpy style cumprod op. Return the cumulative product of the elements along a given axis.
Parameters#
- datarelax.Expr
The input data to the operator.
- axisOptional[int]
Axis along which the cumulative product is computed. The default (None) is to compute the cumprod over the flattened array.
- dtypeOptional[Union[str, DataType]]
Type of the returned array and of the accumulator in which the elements are computed. If dtype is not specified, it defaults to the dtype of data.
- exclusivebool
If false (default), all elements are included in the product. If true, the first element is excluded from the product.
Returns#
- resultrelax.Expr
The result has the same size as data, and the same shape as data if axis is not None. If axis is None, the result is a 1-d array.
Examples#
a = [[1, 2, 3], [4, 5, 6]] cumprod(a) # if axis is not provided, cumprod is done over the flattened input. -> [ 1, 2, 6, 24, 120, 720] cumprod(a, dtype="float32") -> [ 1., 2., 6., 24., 120., 720.] cumprod(a, axis=0) # multiply over rows for each of the 3 columns -> [[1, 2, 3], [4, 10, 18]] cumprod(a, axis=1) -> [[ 1, 2, 6], [ 4, 20, 120]] a = [1, 1, 1, 0, 1, 1, 0] # a is a boolean array cumprod(a, dtype=int32) # dtype should be provided to get the expected results -> [1, 1, 1, 0, 0, 0, 0]
- tvm.relax.op.cumsum(data, axis=None, dtype=None, exclusive=False)[源代码]#
Numpy style cumsum op. Return the cumulative inclusive sum of the elements along a given axis.
Parameters#
- datarelax.Expr
The input data to the operator.
- axisOptional[int]
Axis along which the cumulative sum is computed. The default (None) is to compute the cumsum over the flattened array.
- dtypeOptional[Union[str, DataType]]
Type of the returned array and of the accumulator in which the elements are summed. If dtype is not specified, it defaults to the dtype of data.
- exclusivebool
If false (default), all elements are included in the sum. If true, the first element is excluded from the sum.
Returns#
- resultrelax.Expr
The result has the same size as data, and the same shape as data if axis is not None. If axis is None, the result is a 1-d array.
Examples#
a = [[1, 2, 3], [4, 5, 6]] cumsum(a) # if axis is not provided, cumsum is done over the flattened input. -> [ 1, 3, 6, 10, 15, 21] cumsum(a, dtype="float32") -> [ 1., 3., 6., 10., 15., 21.] cumsum(a, axis=0) # sum over rows for each of the 3 columns -> [[1, 2, 3], [5, 7, 9]] cumsum(a, axis=1) -> [[ 1, 3, 6], [ 4, 9, 15]] a = [1, 0, 1, 0, 1, 1, 0] # a is a boolean array cumsum(a, dtype=int32) # dtype should be provided to get the expected results -> [1, 1, 2, 2, 3, 4, 4]
- tvm.relax.op.dequantize(data, scale, zero_point, axis=-1, out_dtype='float32')[源代码]#
Dequantize op This operator takes input and produces dequantized output. The input tensor can be of any shape. The output shape is the same as input shape.
output = clamp(scale * (input_tensor - zero_point), out_dtype::min, out_dtype::max)
Parameters#
- datatvm.relax.Expr
The input tensor to be dequantized.
- scaletvm.relax.Expr
The input scale.
- zero_pointtvm.relay.Expr
The input zero_point.
- axisint
The channel axis for dequantization. Default value is -1 which corresponds to the last axis.
- out_dtypestr, optional
The data type of the output tensor.
Returns#
- resulttvm.relax.Expr
The computed result.
- tvm.relax.op.divide(x1, x2)[源代码]#
Division with numpy-style broadcasting.
Parameters#
- x1relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.dynamic_strided_slice(x, begin, end, strides)[源代码]#
Dynamic strided slice of a tensor. begin, end, strides can be computed at runtime.
Parameters#
- xExpr
The source tensor to be sliced.
- beginExpr
The indices to begin with in the slicing, inclusive.
- endExpr
The indices indicating end of the slice, exclusive.
- stridesExpr
Specifies the stride values, it can be negative in that case, the input tensor will be reversed in that particular axis. If not specified, it by default is an list of ones of the same length as axes.
Returns#
- retrelax.Expr
The sliced result.
Note#
dyn_strided_slice require the input begin, end and strides to have the same length as rank of data tensor.
- 参数:
x (RelayExpr)
begin (RelayExpr)
end (RelayExpr)
strides (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.einsum(operands, subscripts)[源代码]#
Evaluates the Einstein summation convention on data
Parameters#
- operandsUnion(List[relax.Expr], Tuple[relax.Expr])
A list of expression.
- subscriptsstr
The einsum expression string.
Returns#
- resultrelax.Expr
The output from the einsum op.
- tvm.relax.op.equal(x1, x2)[源代码]#
Broadcasted element-wise test for (lhs == rhs).
Parameters#
- x1relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.erf(x)[源代码]#
Computes the error function of the input.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
Computed error function for each element.
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.ewise_fma(x1, x2, x3)[源代码]#
Elementwise fused multiply-add operator Returns elementwise result of \(x1 * x2 + x3\)
Parameters#
- x1relax.Expr
The left hand operand of the multiplication
- x2relax.Expr
The right hand operand of the multiplication
- x3relax.Expr
The operand of the addition
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
x3 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.exp(x)[源代码]#
Compute element-wise exp of data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.expand_dims(x, axis)[源代码]#
Insert new axes at the positions given by axis.
Parameters#
- xrelax.Expr
The input data to the operator.
- axisUnion[int, List[int]]
The axes at which the input array are expanded. All values are required to lie in range [-data.ndim - 1, data.ndim], with the convention of negative indexing.
Returns#
- resultrelax.Expr
The transformed result.
- tvm.relax.op.eye(n, m=None, k=0, dtype='float32')[源代码]#
Construct a 2-D tensor with ones on the diagonal and zeros elsewhere.
Parameters#
- nUnion[PrimExprLike, PrimValue]
Number of rows in the output.
- mOptional[Union[PrimExprLike, PrimValue]]
Number of columns in the output. If None, defaults to n.
- kUnion[PrimExprLike, PrimValue]
Index of the diagonal: 0 (the default) refers to the main diagonal, a positive value refers to an upper diagonal, and a negative value to a lower diagonal.
- dtypeUnion[str, DataType]
The data type of the created tensor.
Returns#
- resultrelax.Expr
The result tensor.
- tvm.relax.op.eye_like(x, k=0, dtype=None)[源代码]#
Return a 2-D tensor with ones on the diagonal and zeros elsewhere, with the same shape as the input tensor.
Parameters#
- xrelax.Expr
The input tensor, which provides the shape, and dtype when the dtype field is not specified.
- kUnion[PrimExprLike, PrimValue]
Index of the diagonal: 0 (the default) refers to the main diagonal, a positive value refers to an upper diagonal, and a negative value to a lower diagonal.
- dtypeOptional[Union[str, DataType]]
The data type of the created tensor. If dtype is not given, it will by default use the dtype of the input tensor.
Returns#
- resultrelax.Expr
The result tensor.
- tvm.relax.op.flatten(x)[源代码]#
Flatten all the tensor dimensions into one.
Parameters#
- xrelax.Expr
The input data to the operator.
Returns#
- resultrelax.Expr
The flattened result.
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.flip(data, axis)[源代码]#
Reverses the order of elements along given axis while preserving array shape.
Parameters#
- datarelax.Expr
The input data to the operator.
- axis: int
axis to flip on
Returns#
- retrelax.Expr
The computed result.
Examples#
x = [[1., 2.], [3., 4.]] relax.flip(x, axis=0) = [[3., 4.], [1., 2.]] relax.flip(x, axis=1) = [[2., 1.], [4., 3.]]
- tvm.relax.op.floor(x)[源代码]#
Take floor of input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.floor_divide(x1, x2)[源代码]#
Floor division with numpy-style broadcasting.
Parameters#
- x1relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.floor_mod(x1, x2)[源代码]#
Floor modulo with numpy-style broadcasting.
Parameters#
- x1Expr
The first input tensor.
- x2Expr
The second input tensor.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.full(shape, fill_value, dtype=None)[源代码]#
Fill array with scalar value.
Parameters#
- shapeUnion[Tuple[PrimExprLike], Expr]
The shape of the created tensor.
- fill_valuerelax.Expr
The value to fill. Must be a scalar tensor.
- dtypeOptional[Union[str, DataType]]
The data type of the created tensor. If dtype is not given, it will by default use the dtype of fill_value.
Returns#
- resultrelax.Expr
The result tensor.
- tvm.relax.op.full_like(x, fill_value, dtype=None)[源代码]#
Construct a tensor such that - its shape is the same as the input data tensor's shape, - its value is filled with the input scalar fill value.
Parameters#
- xrelax.Expr
The input tensor, which provides the shape, and dtype when the dtype field is not specified.
- fill_valuerelax.Expr
The value to fill. Must be a scalar tensor.
- dtypeOptional[Union[str, DataType]]
The data type of the created tensor. If dtype is not given, it will by default use the dtype of the input tensor.
Returns#
- resultrelax.Expr
The result tensor.
- 参数:
x (RelayExpr)
fill_value (RelayExpr)
dtype (str | DataType | None)
- 返回类型:
RelayExpr
- tvm.relax.op.greater(x1, x2)[源代码]#
Broadcasted element-wise test for (lhs > rhs).
Parameters#
- x1relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.greater_equal(x1, x2)[源代码]#
Broadcasted element-wise test for (lhs >= rhs).
Parameters#
- x1relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.hint_on_device(data, dst_vdevice)[源代码]#
It provides a hint specifying the device on which the input data should be executed. This hint is utilized by RealizeVDevice to propagate the virtual device."
Parameters#
- dataExpr
The tensor to be copied.
- dst_deviceVDevice
The destination device where the data is supposed to be executed.
Returns#
- resultExpr
The result.
- 返回类型:
RelayExpr
- tvm.relax.op.invoke_closure(closure, args, sinfo_args)[源代码]#
Invoke a closure.
Parameters#
- closureExpr
The VMClosure object.
- argsExpr
The input arguments.
- type_args: Union[List[StructInfo], StructInfo]
The structure info arguments of the CallNode
Returns#
- ret: Call
A call to invoke_closure.
- 参数:
closure (RelayExpr)
args (RelayExpr)
sinfo_args (List[StructInfo] | StructInfo)
- 返回类型:
- tvm.relax.op.invoke_pure_closure(closure, args, sinfo_args)[源代码]#
Invoke a closure and indicate to the compiler that it is pure.
Note: This should be used for cases when the user knows that calling the closure with these arguments will in reality not cause any side effects. If it is used for a call that _does_ result in side effects, then the compiler may end up removing, reordering, or repeating that call, with no guarantees made about any side effects from the callee.
Parameters#
- closureExpr
The VMClosure object.
- argsExpr
The input arguments.
- type_args: Union[List[StructInfo], StructInfo]
The structure info arguments of the CallNode
Returns#
- ret: Call
A call to invoke_pure_closure.
- 参数:
closure (RelayExpr)
args (RelayExpr)
sinfo_args (List[StructInfo] | StructInfo)
- 返回类型:
- tvm.relax.op.isfinite(x)[源代码]#
Check if input value is finite.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.isinf(x)[源代码]#
Check if input value is infinite.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.isnan(x)[源代码]#
Check if input value is Nan.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.layout_transform(x, index_map, pad_value=None, axis_separators=None, input_axis_separators=None)[源代码]#
Modifies the layout of a tensor.
Parameters#
- xrelax.Expr
The input tensor to the operator.
- index_mapUnion[Callable, IndexMap]
The transformation to apply.
- pad_valueOptional[Union[int, float, PrimValue]]
The value used for padding if the transformation results in implicit padding. If not specified, any value can be used.
- axis_separatorsOptional[Union[int, IndexMap.AXIS_SEPARATOR]]
The axis_separators for index_map to create non flat buffers.
Returns#
- resultrelax.Expr
The transformed tensor.
- tvm.relax.op.left_shift(x1, x2)[源代码]#
Bitwise Shift Left Parameters ---------- x1 : relax.Expr
The input tensor to be shifted.
- x2relax.Expr
The number of positions to shift.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.less(x1, x2)[源代码]#
Broadcasted element-wise test for (lhs < rhs).
Parameters#
- x1relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.less_equal(x1, x2)[源代码]#
Broadcasted element-wise test for (lhs <= rhs).
Parameters#
- x1relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.linear(data, weight, bias=None, out_dtype=None)[源代码]#
Applies a linear transformation to the incoming data: y = xA^T + b
Parameters#
- datarelax.Expr
The input data.
- weightrelax.Expr
The weight tensor.
- biasOptional[Expr]
The bias tensor.
- out_dtype: Optional[Union[str, DataType]]
The data type of the matmul result. When it is not specified, the output dtype will be the same as input dtype.
Notes#
Relax does not regard the Linear Op as a primitive Op, while combine the transpose, matmul and add op to implement it.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
data (RelayExpr)
weight (RelayExpr)
bias (RelayExpr | None)
out_dtype (str | DataType | None)
- 返回类型:
RelayExpr
- tvm.relax.op.log(x)[源代码]#
Compute element-wise natural logarithm of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.logical_and(x1, x2)[源代码]#
Logical AND Parameters ---------- x1 : relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.logical_not(x)[源代码]#
Compute logical NOT of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.logical_or(x1, x2)[源代码]#
Logical OR Parameters ---------- x1 : relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.logical_xor(x1, x2)[源代码]#
Logical XOR Parameters ---------- x1 : relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.make_closure(func, args)[源代码]#
Create a closure with free variables and return the closure.
Parameters#
- funcExpr
The closure, can be ExternFunc or PrimFunc.
- argsExpr
The input arguments.
Returns#
- ret: Object
The VMClosure.
- 参数:
func (RelayExpr)
args (RelayExpr)
- 返回类型:
Object
- tvm.relax.op.masked_fill(x, mask, value)[源代码]#
Fill a tensor by a specified value in places defined by a mask. Parameters ---------- x : relax.Expr
The input data to the operator.
- maskrelax.Expr
The mask.
- valuerelax.Expr
The value to set in the input tensor.
Returns#
- resultrelax.Expr
The filled tensor.
- 参数:
x (RelayExpr)
mask (RelayExpr)
value (RelayExpr)
- tvm.relax.op.matmul(x1, x2, out_dtype=None)[源代码]#
General matrix multiplication of two tensors, with broadcasting on batched dimensions.
The semantics and output shape deduction rule is specified as https://data-apis.org/array-api/latest/API_specification/generated/array_api.matmul.html.
Parameters#
- x1relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
- out_dtype: Optional[Union[str, DataType]]
The data type of the matmul result. When it is not specified, the output dtype will be the same as input dtype.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
out_dtype (str | DataType | None)
- 返回类型:
RelayExpr
- tvm.relax.op.max(x, axis=None, keepdims=False)[源代码]#
Computes the max of tensor elements over given axes.
Parameters#
- xrelax.Expr
The input data tensor
- axisOptional[Union[int, List[int]]]
Axis or axes along which a max operation is performed. The default, axis=None, will compute the max of all elements in the input tensor. Negative indexing is supported.
- keepdimsbool
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 tensor.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.maximum(x1, x2)[源代码]#
Element-wise maximum
Parameters#
- x1relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.mean(x, axis=None, keepdims=False)[源代码]#
Computes the mean of tensor elements over given axes.
Parameters#
- xrelax.Expr
The input data tensor
- axisOptional[Union[int, List[int]]]
Axis or axes along which a mean operation is performed. The default, axis=None, will compute the mean of all elements in the input tensor. Negative indexing is supported.
- keepdimsbool
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 tensor.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.min(x, axis=None, keepdims=False)[源代码]#
Computes the min of tensor elements over given axes.
Parameters#
- xrelax.Expr
The input data tensor
- axisOptional[Union[int, List[int]]]
Axis or axes along which a min operation is performed. The default, axis=None, will compute the min of all elements in the input tensor. Negative indexing is supported.
- keepdimsbool
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 tensor.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.minimum(x1, x2)[源代码]#
Element-wise minimum
Parameters#
- x1relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.mod(x1, x2)[源代码]#
Modulo with numpy-style broadcasting.
Parameters#
- x1Expr
The first input tensor.
- x2Expr
The second input tensor.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.multinomial_from_uniform(prob, uniform_sample, sample_indices, dtype='int64')[源代码]#
Returns a tensor where each row contains the index sampled from the multinomial probability distribution located in the corresponding row of tensor prob.
Notes#
For better cpu performance, use 'vm.builtin.multinomial_from_uniform'. For accurate results, ensure probabilities are between 0 and 1 and sum to 1.
Parameters#
- probrelax.Expr
A 2-D tensor of shape (batch, vocab_size) representing probability distributions. Each row is a distribution across vocabulary for a batch, where: Values range from [0, 1], indicating the probability of each vocabulary item. The sum of values in each row is 1, forming a valid distribution.
- uniform_samplerelax.Expr
The uniformly sampled 2-D tensor with the shape (n, 1). Values range from 0 to 1, indicating probabilities sampled uniformly.
- sample_indicesrelax.Expr
The 2-D tensor with the shape [n, 1], which indicates the specific probability distribution to sample from. The value of sample_indices[i] determines that the ith token should be sampled from the sample_indices[i]th probability distribution. For instance, if there are 3 distinct probability distributions and the requirement is to sample 2, 3, and 4 tokens from each, then sample_indices would be [0, 0, 1, 1, 1, 2, 2, 2, 2].
- dtypestr
The data type of the output tensor.
Returns#
- resultrelax.Expr
The computed tensor with shape (n, 1).
Examples#
prob = [[0.2, 0.3, 0.5], [0.3, 0.4, 0.3]] usample = [[0.4], [0.9]] sample_indices = [[0], [1]] multinomial_from_uniform(prob, usample) -> [[1], [2]] multinomial_from_uniform(prob, usample, sample_indices) -> [[1], [2]]
- 参数:
prob (RelayExpr)
uniform_sample (RelayExpr)
sample_indices (RelayExpr)
dtype (str)
- 返回类型:
RelayExpr
- tvm.relax.op.multiply(x1, x2)[源代码]#
Multiplication with numpy-style broadcasting.
Parameters#
- x1Expr
The first input tensor.
- x2Expr
The second input tensor.
Returns#
- resultExpr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.negative(x)[源代码]#
Compute element-wise negative of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.nonzero(x)[源代码]#
Find the indices of elements of a tensor that are non-zero.
Parameters#
- xrelax.Expr
The input data tensor.
Returns#
- resultrelax.Expr
A (n+1)-D tensor containing indices of non-zero elements.
Note#
This function is equivalent to onnx.nonzero.
Examples#
x = [[0, 1], [2, 0]] nonzero(x) = [[0, 1], [1, 0]]
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.not_equal(x1, x2)[源代码]#
Broadcasted element-wise test for (lhs != rhs).
Parameters#
- x1relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.null_value()[源代码]#
Create a call node that represents a null value object.
Returns#
- ret: Call
The created call node.
- 返回类型:
- tvm.relax.op.one_hot(indices, on_value, off_value, depth, axis=-1)[源代码]#
Returns a one-hot tensor.
Parameters#
- indicesrelax.Expr
The indices to set to on_value.
- on_valuerelax.PrimValue
The value to fill at indices.
- off_valuerelax.PrimValue
The value to fill at other locations.
- depthint
The depth of the one-hot dimension.
- axisint, optional
The axis to fill. Default is -1 which adds a new dimension at the end.
Returns#
- resultrelax.Expr
The computed result.
Examples#
indices = [0, 1, 2] depth = 3 on_value = 1 off_value = 0 one_hot(indices, on_value, off_value, depth) = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
- tvm.relax.op.ones(shape, dtype)[源代码]#
Construct a tensor of all ones, with the input shape and dtype.
Parameters#
- shapeUnion[Tuple[PrimExprLike], Expr]
The shape of the created tensor.
- dtypeUnion[str, DataType]
The data type of the created tensor.
Returns#
- resultrelax.Expr
The result tensor.
- tvm.relax.op.ones_like(x, dtype=None)[源代码]#
Construct a tensor with all ones, with shape of the input tensor shape.
Parameters#
- xrelax.Expr
The input tensor, which provides the shape, and dtype when the dtype field is not specified.
- dtypeOptional[Union[str, DataType]]
The data type of the created tensor. If dtype is not given, it will by default use the dtype of the input tensor.
Returns#
- resultrelax.Expr
The result tensor.
- 参数:
x (RelayExpr)
dtype (str | DataType | None)
- 返回类型:
RelayExpr
- tvm.relax.op.permute_dims(x, axes=None)[源代码]#
Permutes the dimensions of an array.
Parameters#
- xrelax.Expr
The input data to the operator.
- axesOptional[List[int]]
The target axes order. If not specified, permute_dims will reverse the order of all axes.
Returns#
- resultrelax.Expr
The transposed result.
- tvm.relax.op.power(x1, x2)[源代码]#
Power with numpy-style broadcasting.
Parameters#
- x1relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- tvm.relax.op.print(*values, format='')[源代码]#
Print op to print the values
Parameters#
- valuesList[Expr]
The values to print.
- format: Union[str, Expr]
The format string or StringImm.
Returns#
- resultExpr
A relax Call, which will print the value during runtime.
- tvm.relax.op.prod(x, axis=None, keepdims=False)[源代码]#
Computes the product of tensor elements over given axes.
Parameters#
- xrelax.Expr
The input data tensor
- axisOptional[Union[int, List[int]]]
Axis or axes along which a product is performed. The default, axis=None, will compute the product of all elements of the input tensor. Negative indexing is supported.
- keepdimsbool
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 tensor.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.quantize(data, scale, zero_point, axis=-1, out_dtype='int8')[源代码]#
Quantize op This operator takes input and produces quantized output. The input tensor can be of any shape. The output shape is the same as input shape.
Q_output = clamp((round(input_tensor/scale) + zero_point), out_dtype::min, out_dtype::max)
Parameters#
- datatvm.relax.Expr
The input tensor to be quantized.
- scaletvm.relax.Expr
The output scale.
- zero_pointtvm.relay.Expr
The output zero_point.
- axisint
The channel axis for quantization. Default value is -1 which corresponds to the last axis.
- out_dtypestr, optional
The data type of the output tensor.
Returns#
- resulttvm.relax.Expr
The computed result.
- tvm.relax.op.register_gradient(op_name, fgradient=None, level=10)[源代码]#
Register operator gradient function for a relax operator.
Parameters#
- op_name: str
The name of the op.
- fgradient: function (orig_var: Var, orig_call: Call, output_grad: Var, ctx: BlockBuilder)
-> partials: List[Expr]
The gradient function being used.
- level: int
The priority level
- tvm.relax.op.repeat(data, repeats, axis=None)[源代码]#
Repeats elements of an array.
Parameters#
- datarelax.Expr
The input tensor.
- repeatsint
The number of repetitions.
- axis: Optional[int]
The axis along which to repeat values. The negative numbers are interpreted counting from the backward. By default, use the flattened input array, and return a flat output array.
Returns#
- retrelax.Expr
The computed result.
Examples#
x = R.const([[1, 2], [3, 4]]) lv1 = R.repeat(x, repeats=2) # lv1 == [1, 1, 2, 2, 3, 3, 4, 4] lv2 = R.repeat(x, repeats=2, axis=1) # lv2 == [[1., 1., 2., 2.], # [3., 3., 4., 4.]]
- tvm.relax.op.reshape(x, shape)[源代码]#
Reshape the input array.
-1infers 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.x.shape = (2, 3, 4), shape = (6, 1, -1), result.shape = (6, 1, 4) x.shape = (2, 3, 4), shape = (3, -1, 8), result.shape = (3, 1, 8) x.shape = (2, 3, 4), shape = (-1,), result.shape = (24,)
Parameters#
- xrelax.Expr
The input data to the operator.
- shapeUnion[Tuple[PrimExprLike], Expr]
The new shape. Should be compatible with the original shape.
Returns#
- resultrelax.Expr
The reshaped result.
Note#
The
-1inference is only performed at compile-time. That is to say, in any case the dimension length of-1cannot be inferred in compile-time, an error will be thrown.
- tvm.relax.op.right_shift(x1, x2)[源代码]#
Bitwise Shift Right Parameters ---------- x1 : relax.Expr
The input tensor to be shifted.
- x2relax.Expr
The number of positions to shift.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.round(x)[源代码]#
Rounds each element of the input data to nearest integer.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.rsqrt(x)[源代码]#
Compute element-wise reciprocal square root of the input data.
\[1/sqrt(x)\]Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.scatter_elements(data, indices, updates, axis=0, reduction='update')[源代码]#
ONNX style scatter elements. This operation updates its value in data to values specified by updates at specific index positions specified by indices. For example, in 2D tensor, the update corresponding to the [i][j] entry is performed as below:
output[indices[i][j]][j] = updates[i][j] if axis = 0 output[i][indices[i][j]] = updates[i][j] if axis = 1
When the reduction is set to some reduction function f, the update corresponding to [i][j] entry is performed as below:
output[indices[i][j]][j] += f(output[indices[i][j]][j], updates[i][j]) if axis = 0 output[i][indices[i][j]] += f(output[i][indices[i][j]], updates[i][j]) if axis = 1
Where f is update, add, mul, mean, max, min.
Parameters#
- datarelax.Expr
The input data to the operator.
- indices: relax.Expr
The index positions to update in data.
- updates: relax.Expr
Values to replace to.
- axis: int
Axis to scatter on.
- reduction: str
Type of reduction to apply: update, add, mul, mean, max, min. It is "update" by default.
Returns#
- resultrelax.Expr
The result has the same size as data, and the same shape as data
Examples#
# inputs data = [ [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], ] indices = [ [1, 0, 2], [0, 2, 1], ] updates = [ [1.0, 1.1, 1.2], [2.0, 2.1, 2.2], ] axis = 0 reduction = "update" # output P output = [ [2.0, 1.1, 0.0] [1.0, 0.0, 2.2] [0.0, 2.1, 1.2] ]
- tvm.relax.op.scatter_nd(data, indices, updates, reduction='update')[源代码]#
Scatter updates into an array according to indices.
Parameters#
- data: relax.Expr
The input data to be updated.
- indices: relax.Expr
The index positions to update in data.
- updates: relax.Expr
Values to replace to.
- reduction: str
Type of reduction to apply: update, add, mul, max, min. It is "update" by default.
Returns#
- resultrelax.Expr
The result has the same shape as data.
Examples#
# inputs data = [1, 2, 3, 4, 5, 6, 7, 8] indices = [[4], [3], [1], [7]] updates = [9, 10, 11, 12] # output output = [1, 11, 3, 10, 9, 6, 7, 12]
- 参数:
data (RelayExpr)
indices (RelayExpr)
updates (RelayExpr)
reduction (str)
- 返回类型:
RelayExpr
- tvm.relax.op.shape_of(expr)[源代码]#
Get shape of a tensor.
Parameters#
- exprExpr
The input Expr.
Returns#
- resultExpr
A relax Call, which gets the shape of the input
- 参数:
expr (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.shape_to_tensor(expr)[源代码]#
Convert shape to tensor expr. Parameters ---------- expr : Expr
The input Expr
Returns#
- resultExpr
A relax Call, which transforms the shape values to the tensor
- 参数:
expr (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.sigmoid(x)[源代码]#
Compute element-wise sigmoid of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.sign(x)[源代码]#
Returns an indication of the sign of a number for each element of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.sin(x)[源代码]#
Compute element-wise sin of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.sinh(x)[源代码]#
Compute element-wise sinh of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.sort(x, axis=-1, descending=False)[源代码]#
Performs sorting along the given axis and returns an array in sorted order.
Parameters#
- xrelax.Expr
The input tensor.
- axisint
Axis along which to sort the input tensor. By default the last axis of the input is used.
- descendingbool
Whether to sort in descending order, the default is False
Returns#
- outrelax.Expr
Sorted tensor.
- tvm.relax.op.split(x, 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 possible). Last section will be smaller if the tensor size along the given dimension is not divisible by the integer.
If indices_or_sections is a tuple of mixture of int or PrimExpr, the entries indicate the indices where along axis the array is split.
Parameters#
- xrelax.Expr
The tensor to be split.
- indices_or_sectionsUnion[int, List[PrimExprLike]]
Indices or sections to split into. Accepts an int or a list.
- axisint
The axis over which to split.
Returns#
- retrelax.Expr
The computed result.
- tvm.relax.op.sqrt(x)[源代码]#
Compute element-wise square root of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.square(x)[源代码]#
Squares each element of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.squeeze(x, axis=None)[源代码]#
Squeeze axes in the array.
Parameters#
- xrelax.Expr
The input data to the operator.
- axisOptional[Union[int, 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#
- resultrelax.Expr
The squeezed result.
- tvm.relax.op.std(x, axis=None, keepdims=False)[源代码]#
Computes the standard deviation of tensor elements over given axes.
Parameters#
- xrelax.Expr
The input data tensor
- axisOptional[Union[int, List[int]]]
Axis or axes along which a standard deviation is performed. The default, axis=None, will compute the std of all elements of the input tensor. Negative indexing is supported.
- keepdimsbool
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 tensor.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.strided_slice(x, axes, begin, end, strides=None, assume_inbound=False)[源代码]#
Strided slice of a tensor.
Parameters#
- xrelax.Expr
The source tensor to be sliced.
- axesList[int]
Axes along which slicing is applied.
- beginList[PrimExprLike]
The indices to begin with in the slicing, inclusive.
- endList[PrimExprLike]
The indices indicating end of the slice, exclusive.
- stridesOptional[List[PrimExprLike]]
Specifies the stride values, it can be negative in that case, the input tensor will be reversed in that particular axis. If not specified, it by default is an list of ones of the same length as axes.
- assume_inboundbool
Whether to assume the indices are in bound. If it is set to false, out of bound indices will be clipped to the bound.
Returns#
- retrelax.Expr
The sliced result.
Note#
strided_slice require the input begin, end and strides to have the same length as axes.
- 参数:
x (RelayExpr)
axes (RelayExpr)
begin (RelayExpr)
end (RelayExpr)
strides (RelayExpr | None)
assume_inbound (bool)
- 返回类型:
RelayExpr
- tvm.relax.op.subtract(x1, x2)[源代码]#
Subtraction with numpy-style broadcasting.
Parameters#
- x1relax.Expr
The first input tensor.
- x2relax.Expr
The second input tensor.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.sum(x, axis=None, keepdims=False)[源代码]#
Computes the sum of tensor elements over given axes.
Parameters#
- xrelax.Expr
The input data tensor
- axisOptional[Union[int, List[int]]]
Axis or axes along which a sum is performed. The default, axis=None, will sum all of the elements of the input tensor. Negative indexing is supported.
- keepdimsbool
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 tensor.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.take(x, indices, axis=None)[源代码]#
Take elements from a tensor along an axis. Its semantic is mostly similar to numpy.take (https://numpy.org/doc/stable/reference/generated/numpy.take.html), which can cover torch.take (https://pytorch.org/docs/stable/generated/torch.take.html) and onnx.gather (onnx/onnx).
Parameters#
- xrelax.Expr
The source tensor.
- indicesrelax.Expr
The indices of the values to extract.
- axisOptional[int]
The axis over which to select values. If it is none, the input tensor is required to be one-dimensional.
Returns#
- retrelax.Expr
The taken result.
- 参数:
x (RelayExpr)
indices (RelayExpr)
axis (int | None)
- 返回类型:
RelayExpr
- tvm.relax.op.tan(x)[源代码]#
Compute element-wise tan of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.tanh(x)[源代码]#
Compute element-wise tanh of the input data.
Parameters#
- xrelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.tensor_to_shape(expr)[源代码]#
Convert tensor to shape expr. Parameters ---------- expr : Expr
The input Expr
Returns#
- resultExpr
A relax Call, which transforms the tensor values to the shape
- 参数:
expr (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.tile(data, repeats)[源代码]#
Construct an array by repeating data the number of times given by repeats.
If repeats has length l, and data has dimension d, the result will have dimension of max(l, d).
If d < l, data is promoted to be l-dimensional by prepending new axes. So a shape (3,) Tensor is promoted to (1, 3) for 2-D replication, or shape (1, 1, 3) for 3-D replication. If this is not the desired behavior, promote data to d-dimensions manually before calling this function.
If d > l, reps is promoted to length d by pre-pending 1's to it. Thus for a data of shape (2, 3, 4, 5), a reps of (2, 2) is treated as (1, 1, 2, 2).
Parameters#
- datarelax.Expr
The input data to the operator.
- repeatsUnion[int, Tuple[int], List[int]]
The number of repetitions of data along each axis.
Returns#
- retrelax.Expr
The computed result.
Examples#
x = R.const([[1, 2], [3, 4]]) lv1 = R.tile(x, reps=(2, 3)) # lv1 = [[1., 2., 1., 2., 1., 2.], # [3., 4., 3., 4., 3., 4.], # [1., 2., 1., 2., 1., 2.], # [3., 4., 3., 4., 3., 4.]] lv2 = R.tile(x, reps=2) # lv2 = [[1., 2., 1., 2.], # [3., 4., 3., 4.]]
- tvm.relax.op.to_vdevice(data, dst_vdevice)[源代码]#
Copy data to the destination device. This operator helps data transferring between difference devices for heterogeneous execution.
Parameters#
- dataExpr
The tensor to be copied.
- dst_deviceVDevice
The destination device where the data is copied to.
Returns#
- resultExpr
The copied result.
- 返回类型:
RelayExpr
- tvm.relax.op.topk(data, k=1, axis=-1, ret_type='both', largest=True, dtype='int32')[源代码]#
Get the top k elements in an input tensor along the given axis.
ret_type specifies the return type, can be one of ("both", "values", "indices").
Parameters#
- datarelax.Expr
The input data tensor.
- kint
Number of top elements to select. Return all elements if k < 1.
- axisint
Axis long which to sort the input tensor.
- ret_type: str
The return type [both, values, indices]. "both": return both top k data and indices. "values": return top k data only. "indices": return top k indices only.
- largestbool
Whether to return largest or smallest elements. The k smallest elements are returned if largest is False.
- dtypestr
The data type of the indices output.
Returns#
- outrelax.Expr or List[relax.Expr]
The computed result.
- tvm.relax.op.tril(x, k=0)[源代码]#
Return the lower triangular part of a matrix or a batch of matrices.
Parameters#
- xrelax.Expr
The tensor that tril will be applied to. It is required to have at least two dimensions.
- kint
The index indicating the diagonal above which to zero elements. If k = 0, the diagonal is the main diagonal. If k < 0, the diagonal is below the main diagonal. If k > 0, the diagonal is above the main diagonal.
Returns#
- retrelax.Expr
The result tensor.
- 参数:
x (RelayExpr)
k (int | PrimExpr | RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.triu(x, k=0)[源代码]#
Return the upper triangular part of a matrix or a batch of matrices.
Parameters#
- xrelax.Expr
The tensor that triu will be applied to. It is required to have at least two dimensions.
- kint
The index indicating the diagonal below which to zero elements. If k = 0, the diagonal is the main diagonal. If k < 0, the diagonal is below the main diagonal. If k > 0, the diagonal is above the main diagonal.
Returns#
- retrelax.Expr
The result tensor.
- 参数:
x (RelayExpr)
k (int | PrimExpr | RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.unique(x, sorted=True, return_index=False, return_inverse=False, return_counts=False, axis=None)[源代码]#
Find the unique elements in a given tensor. In addition, it optionally returns - the indices of the input tensor that give the unique values; - the indices of the unique tensor that reconstruct the input tensor; - the number of times each unique value comes up in the input tensor.
Parameters#
- xrelax.Expr
The input tensor.
- sortedUnion[bool, Expr]
Whether to sort the unique elements in ascending order before returning as output.
- return_indexUnion[bool, Expr]
Whether to return an additional tensor with indices for where elements in the unique tensor come from the original input.
- return_inverseUnion[bool, Expr]
Whether to return an additional tensor with indices for where elements in the original input ended up in the returned unique list.
- return_countsUnion[bool, Expr]
Whether to return an additional tensor with counts of each unique elements.
- axisOptional
The dimension to apply unique. If not specified, the unique values of the flattened input are returned.
Returns#
- retrelax.Expr
The created relax call with
- tvm.relax.op.variance(x, axis=None, keepdims=False)[源代码]#
Computes the variance of tensor elements over given axes.
Parameters#
- xrelax.Expr
The input data tensor
- axisOptional[Union[int, List[int]]]
Axis or axes along which a variance operation is performed. The default, axis=None, will compute the variance of all elements in the input tensor. Negative indexing is supported.
- keepdimsbool
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 tensor.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.where(condition, x1, x2)[源代码]#
Selecting elements from either the input tensors depending on the value of the condition.
For a given position, return the corresponding value in x1 if condition is True, and return the corresponding value in x2 otherwise.
Parameters#
- conditionrelax.Expr
When True, yield x1; otherwise, yield x2. Must be broadcasting compatible with x1 and x2. Must have boolean dtype.
- x1relax.Expr
The first input tensor. Must be broadcasting compatible with condition and x2.
- x2relax.Expr
The second input tensor. Must be broadcasting compatible with condition and x1.
Returns#
- resultrelax.Expr
The result tensor.
- 参数:
condition (RelayExpr)
x1 (RelayExpr)
x2 (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.wrap_param(data, dtype='float32')[源代码]#
Cast input tensor which is model param to data type if the dtype of the input data is not the same as the given dtype. Parameters ---------- data : relax.Expr
The input data to the operator.
- dtypeUnion[str, DataType]
The target data type
Returns#
- resultrelax.Expr
The casted result.
- 参数:
data (RelayExpr)
dtype (str | DataType)
- 返回类型:
RelayExpr
- tvm.relax.op.zeros(shape, dtype)[源代码]#
Construct a tensor of all zeros, with the input shape and dtype.
Parameters#
- shapeUnion[Tuple[PrimExprLike], Expr]
The shape of the created tensor.
- dtypeUnion[str, DataType]
The data type of the created tensor.
Returns#
- resultrelax.Expr
The result tensor.
- tvm.relax.op.zeros_like(x, dtype=None)[源代码]#
Construct a tensor with all zeros, with shape of the input tensor shape.
Parameters#
- xrelax.Expr
The input tensor, which provides the shape, and dtype when the dtype field is not specified.
- dtypeOptional[Union[str, DataType]]
The data type of the created tensor. If dtype is not given, it will by default use the dtype of the input tensor.
Returns#
- resultrelax.Expr
The result tensor.
- 参数:
x (RelayExpr)
dtype (str | DataType | None)
- 返回类型:
RelayExpr
tvm.relax.op.nn#
Neural network related operators.
- tvm.relax.op.nn.adaptive_avg_pool1d(data, output_size=None, layout='NCW', out_layout=None)[源代码]#
1D adaptive average pooling operator. This operator is experimental.
This operator takes data as input and does 1D average value calculation across each window represented by W.
In the default case, where the data_layout is NCW a data Tensor with shape (batch_size, in_channels, width), to produce an output Tensor with shape (batch_size, in_channels, output_width).
The pooling kernel and stride sizes are automatically chosen for desired output sizes.
- For output_size:
If this argument is not provided, input height and width will be used as output width.
If a single integer is provided for output_size, the output size is (N x C x output_size) for any input (NCW).
Parameters#
- datarelax.Expr
The input data to the operator.
- output_sizeOptional[Union[int, Tuple[int, int]]]
Output height and width. If not specified, it will be the same as the input height and width. If specified, it is required to have length either 1 or 2.
- layoutstr
Layout of the input.
- out_layoutOptional[str]
Layout of the output. If not specified, it is the same as data_layout
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.nn.adaptive_avg_pool2d(data, output_size=None, layout='NCHW', out_layout=None)[源代码]#
2D adaptive average pooling operator. This operator is experimental.
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 shape (batch_size, in_channels, output_height, output_width).
The pooling kernel and stride sizes are automatically chosen for desired output sizes.
- For output_size:
If this argument is not provided, input height and width will be used as output height and width.
If a single integer is provided for output_size, the output size is (N x C x output_size x output_size) for any input (NCHW).
If a tuple of integers (height, width) are provided for output_size, the output size is (N x C x height x width) for any input (NCHW).
Parameters#
- datarelax.Expr
The input data to the operator.
- output_sizeOptional[Union[int, Tuple[int, int]]]
Output height and width. If not specified, it will be the same as the input height and width. If specified, it is required to have length either 1 or 2.
- layoutstr
Layout of the input.
- out_layoutOptional[str]
Layout of the output. If not specified, it is the same as data_layout
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.nn.adaptive_avg_pool3d(data, output_size=None, layout='NCDHW', out_layout=None)[源代码]#
3D adaptive average pooling operator. This operator is experimental.
This operator takes data as input and does 3D average value calculation across each window represented by WxH.
In the default case, where the data_layout is NCDHW a data Tensor with shape (batch_size, in_channels, depth, height, width), to produce an output Tensor with shape (batch_size, in_channels, output_depth, output_height, output_width).
The pooling kernel and stride sizes are automatically chosen for desired output sizes.
- For output_size:
If this argument is not provided, input depth, height and width will be used as output depth, height and width.
If a single integer is provided for output_size, the output size is (N x C x output_size x output_size x output_size) for any input (NCDHW).
If a tuple of integers (depth, height, width) are provided for output_size, the output size is (N x C x depth x height x width) for any input (NCDHW).
Parameters#
- datarelax.Expr
The input data to the operator.
- output_sizeOptional[Union[int, Tuple[int, int]]]
Output height and width. If not specified, it will be the same as the input height and width. If specified, it is required to have length either 1 or 3.
- layoutstr
Layout of the input.
- out_layoutOptional[str]
Layout of the output. If not specified, it is the same as data_layout
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.nn.attention(query, key, value, bias=None, scale=None, causal_mask=None, window_size=None)[源代码]#
Computes fused multi head attention.
All input tensors are of 4-D tensors with BSNH layout.
\[FMA(Q, K, V) = \text{Softmax}(Q @ K^T) @ V\]备注
The input tensor is required to have float16 dtype
Parameters#
- query: relax.Expr
The input query to the operator. The layout of the input query should be (batch_size, seq_len, num_head, head_dim).
- key: relax.Expr
The input key to the operator. The layout of the input key should be (batch_size, seq_len_kv, num_head, head_dim).
- value: relax.Expr
The input value to the operator. The layout of the input value should be (batch_size, seq_len_kv, num_head, head_dim_v).
- bias: Optional[Expr]
The optional attention bias to the operator. The layout of the attention bias should be a 4-D tensor ending with seq_len_kv, and broadcastable to (batch_size, num_head, seq_len, seq_len_kv).
- scale: Optional[float]
The scale value to be applied to the attention score, by default 1 / sqrt(head_dim).
- causal_mask: Optional[str]
The optional causal mask, i.e. 'TopLeft' and 'BottomRight'. For 'TopLeft', the mask matrix is as np.tril(*, k=0), while for 'BottomRight', the mask matrix is as np.tril(*, k=abs(seq_len - seq_len_kv)) For example, with seq_len = 4, seq_len_kv = 2, mask for 'TopLeft':
[[1, 0], [1, 1], [1, 1], [1, 1]]
mask for 'BottomRight':
[[1, 1], [1, 1], [1, 1], [1, 1]]
with seq_len = 2, seq_len_kv = 4, mask for 'TopLeft':
[[1, 0, 0, 0], [1, 1, 0, 0]]
mask for 'BottomRight':
[[1, 1, 1, 0], [1, 1, 1, 1]]
- window_size: Optional[int]
The size of the window for sliding-window attention.
Returns#
- resultrelax.Expr
The computed result. The layout of the output should be (batch_size, seq_len, num_head, head_dim_v).
- tvm.relax.op.nn.attention_var_len(queries, keys, values, seqstart_q, max_seqlen_q, seqstart_k=None, max_seqlen_k=None, scale=None, causal_mask=None, window_size=None)[源代码]#
Computes fused multi head attention over batched sequences of variable lengths.
Given concatenated inputs and sequence lengths information, this operator computes attention for all sequences more efficiently than calling the normal attention operator for each sequence individually.
Parameters#
- queries: relax.Expr
The input queries concatenated along the second axis. Its shape must be (1, total_seq_len, num_head, head_dim).
- keys: relax.Expr
The input keys concatenated along the second axis. Its shape must be (1, total_seq_len_kv, num_head, head_dim).
- values: relax.Expr
The input values concatenated along the second axis. Its shape must be (1, total_seq_len_kv, num_head, head_dim_v).
- seqstart_q: Optional[Expr]
The cumsum of query sequence lengths, prepended with 0. Its dtype must be int32. For example, if the lengths of the sequences that are batched are [2, 5, 3], this tensor has values [0, 2, 7, 10].
- seqstart_k: Optional[Expr]
The cumsum of key sequence lengths, prepended with 0. By default it is the same as seqstart_q.
- max_seqlen_q: Optional[Expr]
The maximum query sequence length in the batch. It must be int32.
- max_seqlen_k: Optional[Expr]
The maximum key sequence length in the batch. It must be int32. By default it is the same as max_seqlen_q.
- scale: Optional[float]
The scale value to be applied to the attention score, by default 1 / sqrt(head_dim).
- causal_mask: Optional[str]
The optional causal mask, i.e. 'TopLeft' and 'BottomRight'. For 'TopLeft', the mask matrix is as np.tril(*, k=0), while for 'BottomRight', the mask matrix is as np.tril(*, k=abs(seq_len - seq_len_kv)) For example, with seq_len = 4, seq_len_kv = 2, mask for 'TopLeft':
[[1, 0], [1, 1], [1, 1], [1, 1]]
mask for 'BottomRight':
[[1, 1], [1, 1], [1, 1], [1, 1]]
with seq_len = 2, seq_len_kv = 4, mask for 'TopLeft':
[[1, 0, 0, 0], [1, 1, 0, 0]]
mask for 'BottomRight':
[[1, 1, 1, 0], [1, 1, 1, 1]]
- window_size: Optional[int]
The size of the window for sliding-window attention.
Returns#
- resultrelax.Expr
The computed result with shape (1, total_seq_len, num_head, head_dim_v).
- tvm.relax.op.nn.avg_pool1d(data, pool_size=(1,), strides=(1,), padding=(0, 0), dilation=(1,), ceil_mode=False, count_include_pad=False, layout='NCW', out_layout=None)[源代码]#
1D average pooling operator.
This operator takes data as input and does 1D average value calculation with in pool_size sized window by striding defined by stride
In the default case, where the data_layout is NCW a data Tensor with shape (batch_size, channels, width), to produce an output Tensor.
The 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#
- datarelax.Expr
The input data to the operator.
- pool_sizeUnion[int, Tuple[int]]
The size of window for pooling. It is required to have length is 1.
- stridesUnion[int, Tuple[int]]
The strides of pooling. It is required to have length is 1.
- paddingUnion[int, Tuple[int, int]]
The padding for pooling. It is required to have length either 1 or 2.
- dilationUnion[int, Tuple[int]]
The dilation of pooling. It is required to have length is 1.
- ceil_modebool
A boolean indicating if use ceil or floor to compute the output shape. By using ceil, every element in the input tensor will be covered by a sliding window.
- count_include_padbool, optional
To include padding to compute the average.
- layoutstr
Layout of the input.
- out_layoutOptional[str]
Layout of the output. If not specified, it is the same as data_layout
Returns#
- resultExpr
The computed result.
- tvm.relax.op.nn.avg_pool2d(data, pool_size=(1, 1), strides=(1, 1), padding=(0, 0), dilation=(1, 1), ceil_mode=False, count_include_pad=False, layout='NCHW', out_layout=None)[源代码]#
2D average pooling operator.
This operator takes data as input and does 2D avarage 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) = \frac{1}{kh * kw} \sum_{m=0, \ldots, kh-1} \sum_{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#
- datarelax.Expr
The input data to the operator.
- pool_sizeUnion[int, Tuple[int, int]]
The size of window for pooling. It is required to have length either 1 or 2.
- stridesUnion[int, Tuple[int, int]]
The strides of pooling. It is required to have length either 1 or 2.
- paddingUnion[int, Tuple[int, ...]]
The padding for pooling. It is required to have length either 1, 2 or 4.
- dilationUnion[int, Tuple[int, int]]
The dilation of pooling. It is required to have length either 1 or 2.
- ceil_modebool
A boolean indicating if use ceil or floor to compute the output shape. By using ceil, every element in the input tensor will be covered by a sliding window.
- count_include_padbool, optional
To include padding to compute the average.
- layoutstr
Layout of the input.
- out_layoutOptional[str]
Layout of the output. If not specified, it is the same as data_layout
Returns#
- resultExpr
The computed result.
- tvm.relax.op.nn.avg_pool3d(data, pool_size=(1, 1, 1), strides=(1, 1, 1), padding=(0, 0, 0), dilation=(1, 1, 1), ceil_mode=False, count_include_pad=False, layout='NCDHW', out_layout=None)[源代码]#
2D average pooling operator.
This operator takes data as input and does 3D average value calculation with in pool_size sized window by striding defined by stride
In the default case, where the data_layout is NCDHW a data Tensor with shape (batch_size, channels, depth, height, width), to produce an output Tensor.
The 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#
- datarelax.Expr
The input data to the operator.
- pool_sizeUnion[int, Tuple[int, int, int]]
The size of window for pooling. It is required to have length either 1 or 3.
- stridesUnion[int, Tuple[int, int, int]]
The strides of pooling. It is required to have length either 1 or 3.
- paddingUnion[int, Tuple[int, ...]]
The padding for pooling. It is required to have length either 1, 3 or 6.
- dilationUnion[int, Tuple[int, int, int]]
The dilation of pooling. It is required to have length either 1 or 3.
- ceil_modebool
A boolean indicating if use ceil or floor to compute the output shape. By using ceil, every element in the input tensor will be covered by a sliding window.
- count_include_padbool, optional
To include padding to compute the average.
- layoutstr
Layout of the input.
- out_layoutOptional[str]
Layout of the output. If not specified, it is the same as data_layout
Returns#
- resultExpr
The computed result.
- tvm.relax.op.nn.batch_norm(data, gamma, beta, moving_mean, moving_var, axis, epsilon=1e-05, center=True, scale=True, momentum=0.1)[源代码]#
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}\]Both mean and var returns a scalar by treating the input as a vector.
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]\]Assume the input has size k on axis 1, then both
gammaandbetahave shape (k,).Besides the inputs and the outputs, this operator accepts two auxiliary states,
moving_meanandmoving_var, which are k-length vectors. They are global statistics for the whole dataset, which are updated bymoving_mean = moving_mean * momentum + data_mean * (1 - momentum) moving_var = moving_var * momentum + data_var * (1 - momentum)
The parameter
axisspecifies 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.备注
This operator has two modes:
- Training mode.
Use the mean and var computed from THIS batch to normalize.
Update and then return the running mean and running var.
- Inference mode.
Use the running_mean and running_var parameters to normalize.
Do not update the running mean and running var. Just return the original value.
In the legalization stage, this operator will be legalized to the training mode by default.
You can use tvm.relax.transform.DecomposeOpsForInference to decompose the operator, so it executes the inference mode computation. Similarly, use tvm.relax.transform.DecomposeOpsForTraining to execute the training mode computation.
Parameters#
- datarelax.Expr
The input data to the operator.
- gammarelax.Expr
The gamma scale factor.
- betarelax.Expr
The beta offset factor.
- moving_meanrelax.Expr
Running mean of input.
- moving_varrelax.Expr
Running variance of input.
- axisint
The axis along which the normalization is applied.
- epsilonfloat
Small float added to variance to avoid dividing by zero.
- centerbool
Indicating if the beta offset will be added to the normalized tensor.
- scalebool
Indicating if the gamma scale will be multiplied.
- momentumfloat
The value used for the moving_mean and moving_var update.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.nn.conv1d(data, weight, strides=1, padding=0, dilation=1, groups=1, data_layout='NCW', kernel_layout='OIW', out_layout=None, out_dtype=None)[源代码]#
1D convolution.
This operator takes the weight as the 1D convolution kernel and convolves it with data to produce an output.
In the default case, where the data_layout is NCW and kernel_layout is OIW, conv1d takes in a data Tensor with shape (batch_size, in_channels, width), and a weight Tensor with shape (channels, in_channels, kernel_w), where kernel_w is the length of the W kernel dimension, to produce an output Tensor with the following rule:
\[\mbox{out}[b, c, x] = \sum_{dx, k} \mbox{data}[b, k, \mbox{strides} * x + dx] * \mbox{weight}[c, k, 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 (NCW for data and OIW for weight), perform the computation, then convert to the out_layout.
Parameters#
- datarelax.Expr
The input data to the operator.
- weightrelax.Expr
The weight expressions.
- stridesUnion[int, Tuple[int]]
The strides of convolution. It is required to have length 1.
- paddingUnion[int, Tuple[int, ...]]
The padding of convolution on both sides of inputs before convolution. It is required to have length either 1 or 2.
- dilationUnion[int, Tuple[int, int]]
Specifies the dilation rate to be used for dilated convolution. It is required to have length 1.
- groupsint
Number of groups to split the input into for grouped convolution. The number of input and output channels should be divisible by the number of groups.
- data_layoutstr
Layout of the input.
- kernel_layoutstr
Layout of the weight.
- out_layoutOptional[str]
Layout of the output. If not specified, it is the same as data_layout
- out_dtypeOptional[Union[str, DataType]]
Specifies the output data type for mixed precision conv1d.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.nn.conv1d_transpose(data, weight, strides=1, padding=0, output_padding=0, dilation=1, groups=1, data_layout='NCW', kernel_layout='IOW', out_layout=None, out_dtype=None)[源代码]#
1D transposed convolution operator.
This operator can be seen as the gradient operator of conv1d.
The output shape can be explained in the simple case when data_layout == "NCW" and kernel_layout == "IOW". Suppose data has shape (N, in_channel, in_w), weight has shape (in_channel, out_channel, weight_w), we need to assure that in_channel % groups == 0. The shape of the output will be (N, out_channel * groups, out_w), where
out_w = ((in_w - 1) * strides[0] + weight_w - 2 * padding[0] + output_padding[0])
Parameters#
- datarelax.Expr
The input data to the operator.
- weightrelax.Expr
The weight expressions.
- stridesUnion[int, Tuple[int]]
The strides of convolution. It is required to have length 1.
- paddingUnion[int, Tuple[int, ...]]
The padding of convolution on both sides of inputs before convolution. It is required to have length either 1 or 2.
- output_paddingUnion[int, Tuple[int, ...]], optional
Used to disambiguate the output shape.
- dilationUnion[int, Tuple[int]]
Specifies the dilation rate to be used for dilated convolution. It is required to have length either 1.
- groupsint
Number of groups to split the input into for grouped convolution. The number of input and output channels should be divisible by the number of groups.
- data_layoutstr
Layout of the input.
- kernel_layoutstr
Layout of the weight.
- out_layoutOptional[str]
Layout of the output. If not specified, it is the same as data_layout
- out_dtypeOptional[Union[str, DataType]]
Specifies the output data type for mixed precision conv2d.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.nn.conv2d(data, weight, strides=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1, data_layout='NCHW', kernel_layout='OIHW', out_layout=None, out_dtype=None)[源代码]#
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_h, kernel_w), where kernel_h and kernel_w is the lengths of the H and W kernel dimensions, 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#
- datarelax.Expr
The input data to the operator.
- weightrelax.Expr
The weight expressions.
- stridesUnion[int, Tuple[int, int]]
The strides of convolution. It is required to have length either 1 or 2.
- paddingUnion[int, Tuple[int, ...]]
The padding of convolution on both sides of inputs before convolution. It is required to have length either 1, 2 or 4.
- dilationUnion[int, Tuple[int, int]]
Specifies the dilation rate to be used for dilated convolution. It is required to have length either 1 or 2.
- groupsint
Number of groups to split the input into for grouped convolution. The number of input and output channels should be divisible by the number of groups.
- data_layoutstr
Layout of the input.
- kernel_layoutstr
Layout of the weight.
- out_layoutOptional[str]
Layout of the output. If not specified, it is the same as data_layout
- out_dtypeOptional[Union[str, DataType]]
Specifies the output data type for mixed precision conv2d.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.nn.conv2d_transpose(data, weight, strides=(1, 1), padding=(0, 0), output_padding=(0, 0), dilation=(1, 1), groups=1, data_layout='NCHW', kernel_layout='IOHW', out_layout=None, out_dtype=None)[源代码]#
Two dimensional transposed convolution operator.
This operator is intended to be the gradient operator of conv2d. That means, if
out = conv2d(data, weight, strides, padding, dilation),
The gradient w.r.t. data can be calculated as follows:
data_grad = conv2d_transpose(out_grad, weight, strides, padding, output_padding, dilation),
where output_padding is a parameter used to determine the output shape.
The output shape can be explained in the simple case when data_layout == "NCHW" and kernel_layout == "IOHW". Suppose data has shape (N, in_channel, in_h, in_w), weight has shape (in_channel, out_channel, weight_h, weight_w), we need to assure that in_channel % groups == 0. The shape of the output will be (N, out_channel * groups, out_h, out_w), where
out_h = ((in_h - 1) * strides[0] + weight_h - 2 * padding[0] + output_padding[0])
out_w = ((in_w - 1) * strides[1] + weight_w - 2 * padding[1] + output_padding[1])
Parameters#
- datarelax.Expr
The input data to the operator.
- weightrelax.Expr
The weight expressions.
- stridesUnion[int, Tuple[int, int]]
The strides of convolution. It is required to have length either 1 or 2.
- paddingUnion[int, Tuple[int, ...]]
The padding of convolution on both sides of inputs before convolution. It is required to have length either 1, 2 or 4.
- output_paddingUnion[int, Tuple[int, ...]], optional
Used to disambiguate the output shape.
- dilationUnion[int, Tuple[int, int]]
Specifies the dilation rate to be used for dilated convolution. It is required to have length either 1 or 2.
- groupsint
Number of groups to split the input into for grouped convolution. The number of input and output channels should be divisible by the number of groups.
- data_layoutstr
Layout of the input.
- kernel_layoutstr
Layout of the weight.
- out_layoutOptional[str]
Layout of the output. If not specified, it is the same as data_layout
- out_dtypeOptional[Union[str, DataType]]
Specifies the output data type for mixed precision conv2d.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
- 返回类型:
RelayExpr
- tvm.relax.op.nn.conv3d(data, weight, strides=(1, 1, 1), padding=(0, 0, 0), dilation=(1, 1, 1), groups=1, data_layout='NCDHW', kernel_layout='OIDHW', out_layout=None, out_dtype=None)[源代码]#
3D 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 NCDHW and kernel_layout is OIDHW, conv3d takes in a data Tensor with shape (batch_size, in_channels, depth, height, width), and a weight Tensor with shape (channels, in_channels, kernel_d, kernel_h, kernel_w), where kernel_d, kernel_h, and kernel_w are the lengths of the D, H, and W kernel dimensions, to produce an output Tensor with the following rule:
\[\mbox{out}[b, c, z, y, x] = \sum_{dz, dy, dx, k} \mbox{data}[b, k, \mbox{strides}[0] * z + dz, \mbox{strides}[1] * y + dy, \mbox{strides}[2] * x + dx] * \mbox{weight}[c, k, dz, 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 (NCDHW for data and OIDHW for weight), perform the computation, then convert to the out_layout.
Parameters#
- datarelax.Expr
The input data to the operator.
- weightrelax.Expr
The weight expressions.
- stridesUnion[int, Tuple[int, int, int]]
The strides of convolution. It is required to have length either 1 or 3.
- paddingUnion[int, Tuple[int, ...]]
The padding of convolution on both sides of inputs before convolution. It is required to have length either 1, 3 or 6.
- dilationUnion[int, Tuple[int, int, int]]
Specifies the dilation rate to be used for dilated convolution. It is required to have length either 1 or 3.
- groupsint
Number of groups to split the input into for grouped convolution. The number of input and output channels should be divisible by the number of groups.
- data_layoutstr
Layout of the input.
- kernel_layoutstr
Layout of the weight.
- out_layoutOptional[str]
Layout of the output. If not specified, it is the same as data_layout
- out_dtypeOptional[Union[str, DataType]]
Specifies the output data type for mixed precision conv2d.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.nn.cross_entropy_with_logits(predictions, labels)[源代码]#
CrossEntropy with logits between the predictions and labels.
The shape of predictions and labels must be the same. And when ndim >= 2, the first dimension is regarded as the batch_size N. In this case the computed result will divide by N to perform a mean reduction.
\[\text{cross\_entropy\_with\_logits}(x_i, y_i) = \frac{\sum_i -x_i \cdot y_i}{N}\]Parameters#
- predictionsrelax.Expr
The predictions.
- labelsrelax.Expr
The labels (the ground truth values).
Returns#
- resultrelax.Expr
The computed result.
- 参数:
predictions (RelayExpr)
labels (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.nn.dropout(data, rate=0.5)[源代码]#
Applies the dropout operation to the input tensor.
During training, each element of the input is set to zero with probability
p. The whole array is scaled by1/(1-p)to keep the expected sum of the input unchanged.Parameters#
- datarelax.Expr
The input data to the operator.
- ratefloat
The probability for an element to be reset to 0.
Returns#
- resultrelax.Expr
The result of dropout, which is a tuple of two tensors. The first one is the original tensor and the second one is a mask tensor (1.0 where element not dropped, 0.0 where dropped)
- 参数:
data (RelayExpr)
rate (float)
- 返回类型:
RelayExpr
- tvm.relax.op.nn.gelu(data)[源代码]#
Gaussian Error Linear Units function
\[\text{GeLU}(x) = 0.5 * x * (1 + \text{erf}(x * 0.5**0.5))\]where \(erf\) is the Gauss Error function.
Parameters#
- datarelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
data (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.nn.gelu_tanh(data)[源代码]#
Gaussian Error Linear Units function with tanh approximation
\[\text{GELU}(x) = 0.5 * x * (1 + \text{Tanh}(\sqrt(2 / \pi) * (x + 0.044715 * x^3)))\]Parameters#
- datarelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
data (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.nn.group_norm(data, gamma, beta, num_groups, channel_axis, axes, epsilon=1e-05, center=True, scale=True)[源代码]#
Group normalization (Yuxin Wu and et al., 2016). Applies group normalization to the n-dimensional input array. This operator takes an n-dimensional input array. First separate the input array into groups along the channel axis. Then apply layer normalization to each group.
Parameters#
- datarelax.Expr
Input to which group_norm will be applied.
- gammarelax.Expr
The gamma scale factor.
- betarelax.Expr
The beta offset factor.
- num_groupsint
Number of groups to separate the channels into.
- channel_axisint
The index of the channel axis in the input data.
- axesUnion[int, List[int]]
The axes that along which the normalization is applied (excluding the group axis)
- epsilonfloat
Small float added to variance to avoid dividing by zero.
- centerbool
Indicating if the beta offset will be added to the normalized tensor.
- scalebool
Indicating if the gamma scale will be multiplied.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.nn.layer_norm(data, gamma, beta, axes, epsilon=1e-05, center=True, scale=True)[源代码]#
Layer normalization (Lei Ba and et al., 2016). Applies layer normalization to the n-dimensional input array. This operator takes an n-dimensional input array and normalizes the input using the given axis:
\[out = \frac{data - mean(data, axis)}{\sqrt{var(data, axis)+\epsilon}} * gamma + beta\]Unlike batch normalization, the mean and var are computed along the channel dimension.
Assume the input has size k on axis 1, then both gamma and beta have shape (k,).
备注
This operator can be optimized away for inference.
Parameters#
- datarelax.Expr
Input to which layer_norm will be applied.
- gammarelax.Expr
The gamma scale factor.
- betarelax.Expr
The beta offset factor.
- axesUnion[int, List[int]]
The axes that along which the normalization is applied.
- epsilonfloat
Small float added to variance to avoid dividing by zero.
- centerbool
Indicating if the beta offset will be added to the normalized tensor.
- scalebool
Indicating if the gamma scale will be multiplied.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.nn.leakyrelu(data, alpha=0.01)[源代码]#
Rectified linear unit.
\[text{LeakyReLU, negative_slope}(x) = max(x, 0) + negative_slope * min(x, 0)\]Parameters#
- datarelax.Expr
The input data
- alpha: float
Controls the angle of the negative slope, used for nagative inputs. Default value is 0.01
Returns#
- resultrelax.Expr
The computed result.
- 参数:
data (RelayExpr)
alpha (float)
- 返回类型:
RelayExpr
- tvm.relax.op.nn.log_softmax(data, axis=-1)[源代码]#
Computes log softmax.
\[\text{log\_softmax}(x_i) = \log\left( \frac{\exp(x_i)}{\sum_j \exp(x_j)}\right)\]备注
This operator can be optimized away for inference.
Parameters#
- data: relax.Expr
The input data to the operator.
- axis: int
The axis to sum over when computing log softmax. If not specified, it is by default the last axis of the input tensor. Supports negative indexing.
Returns#
- resultrelax.Expr
The computed result.
- 参数:
data (RelayExpr)
axis (int)
- 返回类型:
RelayExpr
- tvm.relax.op.nn.max_pool1d(data, pool_size=(1,), strides=(1,), padding=(0, 0), dilation=(1,), ceil_mode=False, count_include_pad=False, layout='NCW', out_layout=None)[源代码]#
1D maximum pooling operator.
This operator takes data as input and does 1D max value calculation with in pool_size sized window by striding defined by stride.
IIn the default case, where the data_layout is NCW a data Tensor with shape (batch_size, channels, width), to produce an output Tensor.
The 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#
- datarelax.Expr
The input data to the operator.
- pool_sizeUnion[int, Tuple[int, int]]
The size of window for pooling. It is required to have length either 1.
- stridesUnion[int, Tuple[int, int]]
The strides of pooling. It is required to have length either 1.
- paddingUnion[int, Tuple[int, ...]]
The padding for pooling. It is required to have length either 1 or 2.
- dilationUnion[int, Tuple[int, int]]
The dilation of pooling. It is required to have length either 1.
- ceil_modebool
A boolean indicating if use ceil or floor to compute the output shape. By using ceil, every element in the input tensor will be covered by a sliding window.
- count_include_padbool, optional
To include padding to compute the average.
- layoutstr
Layout of the input.
- out_layoutOptional[str]
Layout of the output. If not specified, it is the same as data_layout
Returns#
- resultExpr
The computed result.
- tvm.relax.op.nn.max_pool2d(data, pool_size=(1, 1), strides=(1, 1), padding=(0, 0), dilation=(1, 1), ceil_mode=False, count_include_pad=False, layout='NCHW', out_layout=None)[源代码]#
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#
- datarelax.Expr
The input data to the operator.
- pool_sizeUnion[int, Tuple[int, int]]
The size of window for pooling. It is required to have length either 1 or 2.
- stridesUnion[int, Tuple[int, int]]
The strides of pooling. It is required to have length either 1 or 2.
- paddingUnion[int, Tuple[int, ...]]
The padding for pooling. It is required to have length either 1, 2 or 4.
- dilationUnion[int, Tuple[int, int]]
The dilation of pooling. It is required to have length either 1 or 2.
- ceil_modebool
A boolean indicating if use ceil or floor to compute the output shape. By using ceil, every element in the input tensor will be covered by a sliding window.
- count_include_padbool, optional
To include padding to compute the average.
- layoutstr
Layout of the input.
- out_layoutOptional[str]
Layout of the output. If not specified, it is the same as data_layout
Returns#
- resultExpr
The computed result.
- tvm.relax.op.nn.max_pool3d(data, pool_size=(1, 1, 1), strides=(1, 1, 1), padding=(0, 0, 0), dilation=(1, 1, 1), ceil_mode=False, count_include_pad=False, layout='NCDHW', out_layout=None)[源代码]#
3D maximum pooling operator.
This operator takes data as input and does 3D max value calculation with in pool_size sized window by striding defined by stride.
In the default case, where the data_layout is NCDHW a data Tensor with shape (batch_size, channels, depth, height, width), to produce an output Tensor.
The 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#
- datarelax.Expr
The input data to the operator.
- pool_sizeUnion[int, Tuple[int, int]]
The size of window for pooling. It is required to have length either 1 or 3.
- stridesUnion[int, Tuple[int, int]]
The strides of pooling. It is required to have length either 1 or 3.
- paddingUnion[int, Tuple[int, ...]]
The padding for pooling. It is required to have length either 1, 3 or 6.
- dilationUnion[int, Tuple[int, int]]
The dilation of pooling. It is required to have length either 1 or 3.
- ceil_modebool
A boolean indicating if use ceil or floor to compute the output shape. By using ceil, every element in the input tensor will be covered by a sliding window.
- count_include_padbool, optional
To include padding to compute the average.
- layoutstr
Layout of the input.
- out_layoutOptional[str]
Layout of the output. If not specified, it is the same as data_layout
Returns#
- resultExpr
The computed result.
- tvm.relax.op.nn.nll_loss(predictions, targets, weights=None, reduction='mean', ignore_index=-100)[源代码]#
Negative log likelihood loss.
output[n, i_1, i_2, ..., i_k] = -p * w, where - p = predictions[n, t, i_1, i_2, i_k], - t = targets[n, i_1, i_2, ..., i_k], - w = weights[t] if t != ignore_index else 0
result = reduction(output)
Parameters#
- predictionsrelax.Expr
The predictions. Should be a (k+2)-D Tensor with shape (N, C, d_1, d_2, ..., d_k) where C is the number of target classes.
- targetsrelax.Expr
The target value of each prediction. Should be a (k+1)-D Tensor with shape (N, d_1, d_2, ..., d_k). Must be of int dtype.
- weightsOptional[relax.Expr]
The weight of each target value. Should be a 1-D Tensor with shape (C,). If not specified, it is treated as if having all ones.
- reductionstr
The reduction method to apply to the output. Possible values are "mean", "sum" and "none".
- ignore_indexint
The target value to ignore.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.nn.pad(data, pad_width, pad_value=0, pad_mode='constant')[源代码]#
Padding
This operator takes in a tensor and pads each axis by the specified widths using the specified value.
Parameters#
- data: relax.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
The value used for padding
- pad_mode: 'constant', 'edge', 'reflect'
'constant' pads with constant_value pad_value 'edge' pads using the edge values of the input array 'reflect' pads by reflecting values with respect to the edge
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.nn.relu(data)[源代码]#
Rectified linear unit.
\[\text{ReLU}(x) = \max(x, 0)\]Parameters#
- datarelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
- 参数:
data (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.nn.rms_norm(data, weight, axes=-1, epsilon=1e-05)[源代码]#
Root mean square normalization (Biao Zhang and et al., 2019). Applies root mean square normalization to the n-dimensional input array. This operator takes an n-dimensional input array and normalizes the input using the given axis:
\[out = \frac{data}{\sqrt{mean(data, axis)+\epsilon}} * weight + bias\]Parameters#
- datarelax.Expr
Input to which rms_norm will be applied.
- weightrelax.Expr
The scale factor.
- biasrelax.Expr
The offset factor.
- axesUnion[int, List[int]]
The axes that along which the normalization is applied.
- epsilonfloat
Small float added to square mean to avoid dividing by zero.
Returns#
- resultrelax.Expr
The computed result.
- tvm.relax.op.nn.silu(data)[源代码]#
Sigmoid Linear Unit function
\[\text{SiLU}(x) = x * \text{sigmoid}(x)\]Parameters#
- datarelax.Expr
The input data
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
data (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.nn.softmax(data, axis=-1)[源代码]#
Computes softmax.
\[\text{softmax}(x)_i = \frac{\exp(x_i)}{\sum_j \exp(x_j)}\]Parameters#
- data: relax.Expr
The input data to the operator.
- axis: int
The axis to sum over when computing softmax. If not specified, it is by default the last axis of the input tensor. Supports negative indexing.
Returns#
- resultrelax.Expr
The computed result.
Note#
The input tensor is required to have float dtype
- 参数:
data (RelayExpr)
axis (int)
- 返回类型:
RelayExpr
tvm.relax.op.builtin#
Relax builtin operators.
- tvm.relax.op.builtin.alloc_tensor(shape, dtype, runtime_device_index, storage_scope='global')[源代码]#
Construct a Call to allocate a tensor with specific shape, dtype, runtime_device_index.
Parameters#
- shapeExpr
The shape of the tensor to be allocated.
- dtypeUnion[str, Expr]
The datatype of the tensor to be allocated.
- runtime_device_indexUnion[int, Expr]
The device index indicating on which device the tensor is to be allocated at runtime. Index -1 is reserved for the host device.
- storage_scopeUnion[str, Expr]
The storage scope to allocate the storage to.
Returns#
- resultCall
A relax Call, which gets the allocated tensor.
tvm.relax.op.ccl#
CCL related operators.
- tvm.relax.op.ccl.allgather(x, num_workers, in_group=True)[源代码]#
AllGather operator
Parameters#
- xrelax.Expr
The input tensor.
- num_workerint
The number of workers to gather data from.
- in_groupbool
Whether the gather operation performs globally or in group as default.
Returns#
- resultrelax.Expr
The result of allgather.
- tvm.relax.op.ccl.allreduce(x, op_type='sum', in_group=True)[源代码]#
Allreduce operator
Parameters#
- xrelax.Expr
The input tensor.
- op_typestr
The type of reduction operation to be applied to the input data. Now "sum", "prod", "min", "max" and "avg" are supported.
- in_groupbool
Whether the reduction operation performs globally or in group as default.
Returns#
- resultrelax.Expr
The result of allreduce.
- tvm.relax.op.ccl.broadcast_from_worker0(x)[源代码]#
Broadcast data from worker-0 to all other workers.
Parameters#
- xrelax.Expr
The tensor to be broadcast.
Returns#
- resultrelax.Expr
The same tensor, which has been broadcast to all other workers.
- 参数:
x (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.ccl.scatter_from_worker0(x, num_workers, axis=0)[源代码]#
Perform a scatter operation from worker-0, chunking the given buffer into equal parts.
Parameters#
- xrelax.Expr
The buffer to be divided into equal parts and sent to each worker accordingly.
- num_workerint
The number of workers, i.e. the number of parts the given buffer should be chunked into.
- axisint
The dimension of the tensor to be scattered. Default is 0.
Returns#
- resultrelax.Expr
Chunked Tensor received by different workers.
tvm.relax.op.distributed#
Operators serving for distributed Relax.
- tvm.relax.op.distributed.annotate_sharding(input, device_mesh, placement)[源代码]#
Annotate sharding plan for tensor
Parameters#
- inputrelax.Expr
The input tensor.
- device_mesh: DeviceMesh
The device mesh of the sharding plan
- placement: Placement
The placement of the sharding plan
Returns#
- resultrelax.Expr
The tensor unmodified.
- 参数:
input (RelayExpr)
device_mesh (DeviceMesh)
placement (Placement)
- 返回类型:
RelayExpr
- tvm.relax.op.distributed.call_tir_local_view(gvar, args, out_sinfo, tir_vars=None)[源代码]#
Call a tir.prim_func and return the output. The prim_func should be a worker-local function that is actually executed on each worker, instead of the unpartitioned function. The output of this operator is DTensor or a tuple of DTensors.
Parameters#
- gvarGlobalVar
The GlobalVar referring to a tir PrimFunc.
- argsExpr
The input arguments.
- out_sinfoUnion[DTensorStructInfo, List[DTensorStructInfo]]
The structure info of the call_tir output. It should be a single or a list of DTensorStructInfo. Each one denotes the structure info of a returned tensor.
- tir_varsOptional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]
ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
Returns#
- ret: Call
A call node for the call_tir_local_view operator.
- tvm.relax.op.distributed.redistribute(input, device_mesh, placement)[源代码]#
Redistribute tensor
Parameters#
- inputrelax.Expr
The input tensor.
- device_mesh: DeviceMesh
The device mesh after redistribution
- placement: Placement
The placement after redistribution
Returns#
- resultrelax.Expr
The tensor after redistribution.
- 参数:
input (RelayExpr)
device_mesh (DeviceMesh)
placement (Placement)
- 返回类型:
RelayExpr
- tvm.relax.op.distributed.redistribute_replica_to_shard(input, num_workers, axis)[源代码]#
- Slice tensor into several parts along one axis,
and each worker takes one part. input.struct_info.shape[axis] % num_workers == 0 is required. Each worker must have an identical copy of the input. This is a specialized version of redistribute op.
Parameters#
- inputrelax.Expr
The buffer to be sliced into equal parts.
- num_workerint
The number of workers, i.e. the number of parts the given buffer should be sliced into.
- axisint
The axis of the tensor to be sliced.
Returns#
- resultrelax.Expr
Sliced Tensor kept by each device.
tvm.relax.op.grad#
Operators serving for finding gradient of relax operators.
- tvm.relax.op.grad.Expr#
RelayExpr的别名
- tvm.relax.op.grad.avg_pool2d_backward(output_grad, data, pool_size=(1, 1), strides=(1, 1), padding=(0, 0, 0, 0), dilation=(1, 1), ceil_mode=False, count_include_pad=False, layout='NCHW', out_layout=None)[源代码]#
Backward operator of relax.nn.avg_pool2d. All parameters except output_grad is the same as relax.nn.avg_pool2d. Returns the gradient w.r.t. data.
Parameters#
- output_gradrelax.Expr
The gradient w.r.t. the result of avg_pool2d.
Returns#
- resultrelax.Expr
The gradient w.r.t. data.
- tvm.relax.op.grad.end_checkpoint(input)[源代码]#
Mark the end of checkpoint stage. See tvm.relax.op.grad.start_checkpoint.
Parameters#
- inputrelax.Expr
The output of the checkpoint stage.
Returns#
- resultrelax.Expr
The same tensor as the input.
- 参数:
input (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.grad.max_pool2d_backward(output_grad, data, pool_size=(1, 1), strides=(1, 1), padding=(0, 0, 0, 0), dilation=(1, 1), ceil_mode=False, count_include_pad=False, layout='NCHW', out_layout=None)[源代码]#
Backward operator of relax.nn.max_pool2d. All parameters except output_grad is the same as relax.nn.max_pool2d. Returns the gradient w.r.t. data.
Parameters#
- output_gradrelax.Expr
The gradient w.r.t. the result of max_pool2d.
Returns#
- resultrelax.Expr
The gradient w.r.t. data.
- tvm.relax.op.grad.nll_loss_backward(output_grad, predictions, targets, weights=None, reduction='mean', ignore_index=-100)[源代码]#
Backward operator of relax.nn.nll_loss. All parameters except output_grad is the same as relax.nn.nll_loss. Returns the gradient w.r.t. predictions.
Parameters#
- output_gradrelax.Expr
The gradient w.r.t. the result of nll_loss.
Returns#
- resultrelax.Expr
The gradient w.r.t. predictions.
- tvm.relax.op.grad.no_grad(input)[源代码]#
No gradient dummy operator w.r.t. the input.
Parameters#
- inputrelax.Expr
The corresponding input tensor.
Returns#
- resultrelax.Expr
The no-gradient representation w.r.t. input.
- 参数:
input (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.grad.start_checkpoint(input)[源代码]#
Mark the start of the checkpoint stage. The computation between start_checkpoint and end_checkpoint will be marked as the checkpoint stage.
Rather than storing all intermediate activations of the entire computation graph for computing backward, the checkpointed stage does not save intermediate activations, and instead recomputes them in backward process.
For instance,
` a = relax.Var("a", relax.TensorStructInfo((2, 2), "float32")) b = relax.Var("b", relax.TensorStructInfo((2, 2), "float32")) c = a * 2 d = b * 2 c_cp = start_checkpoint(c) d_cp = start_checkpoint(d) e = c_cp + d_cp e_out = end_checkpoint(e) `Then e will be recomputed in the backward stage.See tvm.relax.transform.Gradient, tvm.relax.testing.nn.checkpoint, tvm.relax.op.grad.end_checkpoint for more information.
Parameters#
- inputrelax.Expr
The tensor marking the input of the checkpoint stage.
Returns#
- resultrelax.Expr
The same tensor as the input.
- 参数:
input (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.grad.take_backward(output_grad, x, indices, axis=None)[源代码]#
Backward operator of relax.take. All parameters except output_grad is the same as relax.take. Returns the gradient w.r.t. x.
Parameters#
- output_gradrelax.Expr
The gradient w.r.t. the result of take.
Returns#
- resultrelax.Expr
The gradient w.r.t. x.
- 参数:
output_grad (RelayExpr)
x (RelayExpr)
indices (RelayExpr)
axis (int | None)
- 返回类型:
RelayExpr
tvm.relax.op.image#
Image operators.
- tvm.relax.op.image.resize2d(data, size, roi=None, layout='NCHW', method='linear', coordinate_transformation_mode='half_pixel', rounding_method='round', cubic_alpha=-0.5, cubic_exclude=0, extrapolation_value=0.0, out_dtype=None)[源代码]#
Image resize2d operator.
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, size[0], size[1])
method indicates the algorithm to be used while calculating the out value and method can be one of ("linear", "nearest_neighbor", "cubic")
Parameters#
- datarelax.Expr
The input data to the operator.
- size: Union[Expr, PrimExprLike, Tuple[PrimExprLike]]
The out size to which the image will be resized. If specified as a list, it is required to have length either 1 or 2. If specified as an Expr, it is required to have ndim 2.
- roi: Optional[Union[float, Tuple[float]]]
The region of interest for cropping the input image. Expected to be of size 4, and format [start_h, start_w, end_h, end_w]. Only used if coordinate_transformation_mode is tf_crop_and_resize.
- layoutstr
Layout of the input.
- methodstr
Scale method to used [nearest_neighbor, linear, cubic].
- coordinate_transformation_modestr
Describes how to transform the coordinate in the resized tensor to the coordinate in the original tensor. Definitions can be found in topi/image/resize.py. [half_pixel, align_corners, asymmetric, pytorch_half_pixel, tf_half_pixel_for_nn, and tf_crop_and_resize].
- rounding_method: str
indicates how to find the "nearest" pixel in nearest_neighbor method [round, floor, ceil]
- cubic_alpha: float
Spline Coefficient for bicubic interpolation
- cubic_exclude: int
Flag to exclude exterior of the image during bicubic interpolation
- extrapolation_value: float
Fill value to use when roi is outside of the image
- out_dtypeOptional[Union[str, DataType]]
The dtype of the output tensor. It it is not specified, the output will have the same dtype as input if not specified.
Returns#
- result: relax.Expr
The resized result.
- 参数:
- 返回类型:
RelayExpr
tvm.relax.op.memory#
Relax memory primitives.
- tvm.relax.op.memory.alloc_storage(size, virtual_device_index, storage_scope, dtype)[源代码]#
Construct a Call to allocate a storage with specific size, virtual_device_index, storage_scope and dtype.
Parameters#
- sizeExpr
The size of the storage to be allocated.
- virtual_device_indexUnion[int, Expr]
The virtual device index indicating on which device the storage is to be allocated. Index -1 is reserved for the host device.
- storage_scopeUnion[str, Expr]
The storage scope to allocate the storage to.
- dtypeUnion[str, Expr]
The datatype of the storage to be allocated.
Returns#
- resultCall
A relax Call, which gets the allocated storage.
- tvm.relax.op.memory.alloc_tensor(storage, offset, shape, dtype)[源代码]#
Construct a Call to allocate a tensor on a certain storage starting from the given offset.
Parameters#
- storageExpr
The storage to allocate the tensor to.
- offsetUnion[int, Expr]
The storage offset to allocate the tensor.
- shapeExpr
The shape of the tensor to be allocated.
- dtypeUnion[str, Expr]
The datatype of the tensor to be allocated.
Returns#
- resultCall
A relax Call, which gets the allocated tensor.
- tvm.relax.op.memory.ensure_zero_offset(data)[源代码]#
Ensure the tensor has elem_offset == 0. A copy will be made if necessary.
Parameters#
- datarelax.Expr
The input tensor
Results#
- resultrelax.Expr
The tensor with elem_offset == 0
- 参数:
data (RelayExpr)
- 返回类型:
RelayExpr
- tvm.relax.op.memory.kill_storage(storage)[源代码]#
Construct a Call to kill a storage.
Parameters#
- storageExpr
The storage to be killed.
Returns#
- resultCall
A relax Call to kill a storage.
- 参数:
storage (RelayExpr)
- 返回类型:
- tvm.relax.op.memory.kill_tensor(tensor)[源代码]#
Construct a Call to kill a tensor.
Parameters#
- tensorExpr
The tensor to be killed.
Returns#
- resultCall
A relax Call to kill a tensor.
- 参数:
tensor (RelayExpr)
- 返回类型:
- tvm.relax.op.memory.view(data, shape=None, dtype=None, relative_byte_offset=None)[源代码]#
Provide a view into an existing tensor
The view may have a different shape, may be a different datatype, and may start at an offset relative to the source array.
Regardless of which combination of these options are used, the view may never access memory that was not accessible through the input data array. This restriction applies even if the data array is itself a view into a shared backing array.
Parameters#
data : relax.Expr
The input data to the operator.
shape : Optional[Union[Sequence[PrimExprLike], Expr]]
The target shape. Should be a relax.ShapeExpr, or a collection that can be converted to a relax.ShapeExpr.
dtype : Optional[Expr]
The target datatype. Should be a relax.ShapeExpr, or a collection that can be converted to a relax.ShapeExpr.
relative_byte_offset: Optional[Expr]
The offset of the output NDArray, relative to the byte offset of data. If None, the offset of the view is the same as the offset of data.
Returns#
- resultrelax.Expr
The tensor view
tvm.relax.op.op_attrs#
The attributes node used for Relax operators
- class tvm.relax.op.op_attrs.CallTIRWithGradAttrs[源代码]#
Attributes used in call_tir_with_grad operator