tvm.script

Examples

bb = relax.BlockBuilder()
a = relax.Var("a", relax.TensorType(shape=(2, 3), dtype="float32"))
b = relax.Var("b", relax.TensorType(shape=(2, 1), dtype="float32"))
c = bb.normalize(relax.op.add(a, b))  # c has TensorType(shape=(2, 3), dtype="float32")

tvm.script.relax.and_(*values: Any) → Any#: Implements tvm.script.ir_builder.parser_protocol.and_().

Construct a tensor with evenly spaced elements.

Parameters:

start (PrimExprLike) – The start of the interval.
end (Optional[PrimExprLike]) – The end of the interval. If not given, it will be set to start, and start will be set to 0.
step (PrimExprLike) – The step size.
dtype (Optional[str | DataType | PrimType]) – The data type of the created tensor.

Returns:

result – The result tensor.

Return type:

tvm.script.relax.arg_(name: str, ty: Any, *, span: SpanEntry | Span | None = None) → Var#

Declare a Relax parameter using the shared argument-hook contract.

Parameters:

name (str) – Source parameter name.
ty (Type, Var, or callable) – Parameter annotation; corresponds to annotation in the shared tvm.script.ir_builder.parser_protocol.arg_() contract. Primitive annotations create a fresh parameter; an existing variable retains its identity.
span (SpanEntry, Span or None, optional) – Source location attached to the constructed IR; None leaves it unspecified.

Returns:

result – The parameter registered in the current function frame.

Return type:

tvm.script.relax.argmax(x: Expr, axis: int | None = None, keepdims: bool = False, *, ty=None, span=None) → Expr#

Computes the argmax of tensor elements over given axis.

Parameters:

x (relax.Expr) – The input data tensor
axis (Optional[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.
keepdims (bool) – 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:

result – The computed result.

Return type:

tvm.script.relax.argmin(x: Expr, axis: int | None = None, keepdims: bool = False, *, ty=None, span=None) → Expr#

Computes the argmin of tensor elements over given axis.

Parameters:

x (relax.Expr) – The input data tensor
axis (Optional[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.
keepdims (bool) – 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:

result – The computed result.

Return type:

tvm.script.relax.argsort(data: Expr, axis: int = -1, descending: bool = False, dtype: str = 'int32', *, ty=None, span=None)#

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:

data (relax.Expr) – The input data tensor.
axis (int) – Axis long which to sort the input tensor.
descending (bool) – Whether to sort in descending order, the default is False
dtype (str) – The data type of the output indices.

Returns:

out – Tensor with same shape as data.

Return type:

tvm.script.relax.asin(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise arc sin of the input data.

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

Note

The input tensor is required to have float dtype

tvm.script.relax.asinh(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise arc sinh of the input data.

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

Note

The input tensor is required to have float dtype

tvm.script.relax.assert_(condition: Any, message: str | tuple[str, Sequence[Any]] | Sequence[Any] = '', *, span: SpanEntry | Span | None = None) → None#: Implements tvm.script.ir_builder.parser_protocol.assert_().

tvm.script.relax.assert_op(condition: Expr, format: str | Expr = '', *values: Expr, ty=None, span=None) → Expr#

Create a call to Relax’s assert_op operation (assert is reserved in Python, so the name must be distinct).

Parameters:

condition (Expr) – The assertion condition.
format (Union[str, Expr]) – The format string or StringImm for the error message. If empty, the values are rendered as a comma-separated list.
values (Expr) – Values used to format the error message if the condition fails. A tuple-valued expression is one value unless explicitly expanded.

Returns:

result – A Call to the Relax assert operation.

Return type:

tvm.script.relax.astype(x: Expr, dtype: str | dtype | PrimType, *, ty=None, span=None) → Expr#

Cast input tensor to the given data type.

Parameters:

x (relax.Expr) – The input data to the operator.
dtype (Union[str, DataType]) – The target data type

Returns:

result – The casted result.

Return type:

tvm.script.relax.atan(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise arc tan of the input data.

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

Note

The input tensor is required to have float dtype

tvm.script.relax.atan2(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Atan2 with numpy-style broadcasting.

Parameters:

x1 (relax.Expr) – The first input tensor (y-coordinates).
x2 (relax.Expr) – The second input tensor (x-coordinates).

Returns:

result – The computed result.

Return type:

tvm.script.relax.atanh(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise arc tanh of the input data.

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

Note

The input tensor is required to have float dtype

tvm.script.relax.bind_(value: ~typing.Any = <MISSING>, *, ty: ~typing.Any | None = None, name: str | None = None, span: ~tvm.script.ir_builder.base.SpanEntry | ~tvm.ir.base.Span | None = None, value_span: ~tvm.script.ir_builder.base.SpanEntry | ~tvm.ir.base.Span | None = None, name_span: ~tvm.script.ir_builder.base.SpanEntry | ~tvm.ir.base.Span | None = None, frame_value: bool = False) → Any#

Implements tvm.script.ir_builder.parser_protocol.bind_().

Relax values emit a binding or MatchCast; value_span belongs to the actual RHS and span to the binding target. Metadata passes through unchanged.

tvm.script.relax.bitwise_and(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Bitwise AND :param x1: The first input tensor. :type x1: relax.Expr :param x2: The second input tensor. :type x2: relax.Expr

Returns:: result – The computed result.
Return type:: relax.Expr

tvm.script.relax.bitwise_not(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute bitwise NOT of the input data.

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

tvm.script.relax.bitwise_or(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Bitwise OR :param x1: The first input tensor. :type x1: relax.Expr :param x2: The second input tensor. :type x2: relax.Expr

Returns:: result – The computed result.
Return type:: relax.Expr

tvm.script.relax.bitwise_xor(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Bitwise XOR :param x1: The first input tensor. :type x1: relax.Expr :param x2: The second input tensor. :type x2: relax.Expr

Returns:: result – The computed result.
Return type:: relax.Expr

tvm.script.relax.break_(*, span: SpanEntry | Span | None = None) → NoReturn#

Implements tvm.script.ir_builder.parser_protocol.break_().

Relax rejects loop-control statements.

tvm.script.relax.broadcast_to(x: Expr, shape: tuple[int | Expr] | Expr, *, ty=None, span=None) → Expr#

Broadcasts a tensor to a specified shape.

Parameters:

x (relax.Expr) – The input data to the operator.
shape (Union[Tuple[PrimExprLike], Expr]) – The target shape.

Returns:

result – The broadcasted tensor.

Return type:

tvm.script.relax.bucketize(input_tensor, boundaries, out_int32=False, right=False, *, ty=None, span=None)#

Returns the indices of the buckets to which each value in the input belongs.

Parameters:

input_tensor (relax.Expr) – N-D tensor containing the search values.
boundaries (relax.Expr) – 1-D tensor, must contain a strictly increasing sequence, or the return value is undefined.
out_int32 (Optional[bool]) – Indicate the output data type. int32 if True, int64 otherwise. Default=False
right (Optional[bool]) – Determines the behavior for values in boundaries. Similar to torch.bucketize

Returns:

result – The computed result with same shape as input_tensor.

Return type:

tvm.script.relax.call_builtin_with_ctx(func: str | Expr, args: Expr, *, ty_args: Type | list[Type] | None = None, ty=None, span=None) → Call#

Call a builtin function func.

Parameters:

func (Expr) – The builtin function to be called.
args (Expr) – The input arguments.
ty_args (Optional[Union[Type, List[Type]]]) – The type arguments to the call node.

Returns:

ret – The created call node.

Return type:

tvm.script.relax.call_dps_packed(func, args, *, ty_args, attrs=None, ty=None, span=None, **kwargs) → Call#

Call a destination-passing packed function and allocate its outputs.

Python string callees become ExternFunc; explicit Expr callees retain their identity. ty_args contains one output type, including a TupleType for multiple results.

The function must be pure apart from writing its designated outputs. Other effects may be removed, reordered or repeated by the compiler.

tvm.script.relax.call_global_var_(function: GlobalVar, args: Sequence[Any]) → Expr#: Implements tvm.script.ir_builder.parser_protocol.call_global_var_().

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:

func (Union[str, ExternFunc]) – The name (global symbol) for a PackedFunc or an ExternFunc node.
args (Expr) – The arguments for the PackedFunc.
inplace_indices (Union[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.
ty_args (Union[Type, List[Type]]) – The list of type information arguments (giving the type information for the returned value).

Returns:

result – A Relax call, corresponding to call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), ty_args)

Return type:

tvm.script.relax.call_packed(func: str, *args: Expr, ty_args: Type | list[Type] | None = None, **kwargs: Any) → Call#

Create a relax Call, which calls a packed function. :param func: The name of extern function. :type func: str :param *args: The arguments. :type *args: Expr :param ty_args: The list of type information arguments. :type ty_args: Optional[Union[Type, List[Type]]] :param kwargs: The keyword arguments. :type kwargs: Expr

Returns:: call – The created Relax Call
Return type:: Call

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:

func (Union[str, ExternFunc, Op]) – The name (global symbol) for a PackedFunc or an ExternFunc node. The explicit relax.call_tir_packed Op is also accepted; its native callee and argument tuple follow as arguments to this wrapper.
args (Expr) – The arguments for the PackedFunc.
ty_args (Union[Type, List[Type]]) – The list of type information arguments (giving the type information for the returned value). Omit this for the native bridge, whose result follows the native signature.

Returns:

result – A Relax call, corresponding to call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), ty_args)

Return type:

tvm.script.relax.call_py_func(func_name: str | Expr, args: Expr, *, ty_args, ty=None, span=None) → Call#

Call a Python function using canonical operands and one output type argument.

func_name names a function in the IRModule’s pyfuncs attribute. args accepts an Expr or a Python tuple through shared Expr conversion. ty_args contains exactly one result type, including a TupleType for tuple-valued results. Omitted ty uses the registered result inference.

tvm.script.relax.call_tir(func, args, *, ty_args, attrs=None, ty=None, span=None, **kwargs) → Call#

Call a destination-passing TIR function and allocate its output tensors.

func is the function’s GlobalVar and args contains its ordered inputs. ty_args contains one output type: a TensorType or a TupleType of tensor results. Omitted ty uses the registered result inference; explicit types and spans are forwarded unchanged.

tvm.script.relax.call_tir_inplace(func, args, *, ty_args, attrs=None, ty=None, span=None, **kwargs) → Call#

Call a TIR function whose selected outputs alias its input tensors.

ty_args contains one output type. In the inplace_indices attribute, entry i >= 0 makes the corresponding output alias input i; -1 allocates a fresh output. At least one output must alias an input.

Although classified as pure, this operation mutates the selected inputs. Only use it after proving there are no live uses or aliases that could observe those mutations. Direct construction is intended for testing; optimization passes normally establish these preconditions.

tvm.script.relax.call_tir_packed(gvar: GlobalVar, args: Expr, *, ty=None, span=None) → Call#

Call a TIRx Function through its native packed-call contract.

Every native parameter is supplied explicitly, in order. Unlike call_tir(), this operator does not allocate destination tensors or interpret destination parameters as results. It performs no implicit copies, casts, layout conversions, device transfers, or redistribution.

The result follows the native declared return type: a supported scalar keeps its exact primitive type, void becomes the empty tuple, and a pointer becomes an Any carrying an opaque pointer. A pointer result does not imply tensor ownership or a lifetime guarantee. Native tensor, nonempty tuple, callable, vector, and other unsupported return types are rejected. There is no out_ty argument.

Parameters:

gvar (GlobalVar) – The GlobalVar referring to a TIRx Function with its native tvm.ir.FuncType signature.
args (Expr) –
The ordered arguments, supplied as an inline Relax tuple, a Python tuple or list, or a single expression. Tensor parameters accept compatible Relax tensors, with dtype, rank, and known shapes checked against the native signature. Remaining supported native tensor constraints are checked by the packed ABI at runtime. Specialized global storage scopes require matching VDevice.memory_scope; non-default layouts, allocated-address contracts, unsupported storage scopes, and unlowered distributed tensors are rejected.

Scalar parameters require the exact primitive dtype: scalar bool, signed or unsigned integers up to 64 bits, or float16/32/64. These same scalar types are supported as direct results. The existing packed integer carrier is signed 64-bit: uint64 values must be in [0, 2**63 - 1]; larger unsigned values are not representable.

Pointer parameters accept Any or a handle-compatible object, including a runtime tensor. A tensor passed to a pointer parameter supplies its DLTensor header handle, not its data pointer. This erased carrier does not prove pointee type, address space, ownership, or lifetime compatibility. Known scalar values cannot serve as pointer arguments. The runtime carrier must satisfy the existing packed ABI’s null, opaque-pointer, DLTensor-pointer, or object-handle check.
to (The call is effectful and may mutate its arguments. For a call known)
effects (have no observable)
wrapper (use an explicit call_pure_packed)
inferred (around the relax.call_tir_packed operator. Purity is never)
ABI. (from the native signature or packed)

Returns:

ret – A call whose result type is derived from the native declared return during Relax type inference.

Return type:

Examples

A native (int64, int64) -> int64 function returns its scalar directly:

result = relax.call_tir_packed(add_scalar, (a, b))

A caller that knows the scalar function has no effects may assert purity:

result = relax.call_pure_packed(
    tvm.ir.Op.get("relax.call_tir_packed"), add_scalar, (a, b)
)

A native copy function with a void return writes to a caller-owned tensor:

relax.call_tir_packed(copy, (source, destination))

tvm.script.relax.call_tir_with_grad(func, args, *, ty_args, attrs=None, ty=None, span=None, **kwargs) → Call#

Call a TIR function with a registered TE gradient rule.

ty_args contains the single output type, including a TupleType for multiple results. The te_grad_name and optional te_grad_kwargs attribute keywords select the rule used by the Gradient pass. Attributes may instead be supplied through attrs.

tvm.script.relax.ceil(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Take ceil of input data.

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

tvm.script.relax.check_well_formed_(function: Function) → None#

Validate a completed Relax function.

Parameters:: function (tvm.relax.Function) – Completed function to validate, without changing its IR.
Raises:: ValueError – If the function fails language variant validation.

Notes

See tvm.script.ir_builder.parser_protocol.check_well_formed_() for the shared whole-module validation coordinator.

Validate the completed Relax function. Whole-module cross-function validation is registered separately with the root coordinator.

tvm.script.relax.clip(x: Expr, min: Expr, max: Expr, *, ty=None, span=None) → Expr#

Clips tensor values to a specified min and max.

Parameters:

x (relax.Expr) – The input data
min (relax.Expr) – The minimum value
max (relax.Expr) – The maximum value

Returns:

result – The computed result.

Return type:

tvm.script.relax.collapse_sum_like(data: Expr, collapse_target: Expr, *, ty=None, span=None) → Expr#

Return a summation of data to the shape of collapse_target.

For details, please see relax.op.collapse_sum_to.

Parameters:

data (relax.Expr) – The input tensor.
collapse_target (relax.Expr) – The tensor whose shape is the shape to collapse to.

Returns:

result – The result tensor after summation.

Return type:

tvm.script.relax.collapse_sum_to(data: Expr, shape: tuple[int | Expr] | Expr, *, ty=None, span=None) → Expr#

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:

data (relax.Expr) – The input tensor.
shape (Union[Tuple[PrimExprLike], relax.Expr]) – The shape to collapse to.

Returns:

result – The result tensor of the given shape after summation.

Return type:

tvm.script.relax.concat(tensors: Expr | list[Expr], axis: int | None = 0, *, ty=None, span=None) → Expr#

Concatenate the input tensors along the given axis.

Parameters:

tensors (Union[relax.Expr, List[relax.Expr]]) – An Expr in Tuple type, containing the tensors to be concatenated, or a list of Tensors.
axis (Optional[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 – The concatenated tensor.

Return type:

Create a constant value.

Parameters:

value (bool | int | float | numpy.ndarray | tvm.runtime.Tensor) – The constant value.
dtype (Optional[str]) – The data type of the resulting constant.

Note

When dtype is None, we use the following rule:

int maps to “int32”
float maps to “float32”
bool maps to “bool”
other using the same default rule as numpy.

tvm.script.relax.constexpr(value: object) → NoReturn#

Mark host control syntax or a JIT specialization annotation.

Parameters:: value (object) – Source expression to evaluate with ordinary Python semantics in a supported control-flow position. The marker itself can also appear as an annotation identifying a JIT specialization parameter.
Raises:: TypeError – If invoked directly instead of being recognized in parsed source.

Notes

The parser recognizes this marker through its fixed namespace path and removes it before execution. Host operators retain ordinary Python behavior. Direct invocation raises TypeError; no builder frame or IR is constructed.

# Source
if I.constexpr(enabled):
    T.evaluate(1)
# Builder
if enabled:
    X.emit_(X.evaluate(1))

tvm.script.relax.continue_(*, span: SpanEntry | Span | None = None) → NoReturn#

Implements tvm.script.ir_builder.parser_protocol.continue_().

Relax rejects loop-control statements.

tvm.script.relax.cos(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise cos of the input data.

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

Note

The input tensor is required to have float dtype

tvm.script.relax.cosh(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise cosh of the input data.

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

Note

The input tensor is required to have float dtype

tvm.script.relax.cpu(dev_id=0)#

Construct a CPU device

Parameters:: dev_id (int, optional) – The integer device id
Returns:: dev – The created device
Return type:: Device

tvm.script.relax.cuda(dev_id=0)#

Construct a CUDA GPU device

Parameters:: dev_id (int, optional) – The integer device id
Returns:: dev – The created device
Return type:: Device

tvm.script.relax.cumprod(data: Expr, axis: int | None = None, dtype: str | dtype | None = None, exclusive: bool = False, *, ty=None, span=None)#

Numpy style cumprod op. Return the cumulative product of the elements along a given axis.

Parameters:

data (relax.Expr) – The input data to the operator.
axis (Optional[int]) – Axis along which the cumulative product is computed. The default (None) is to compute the cumprod over the flattened array.
dtype (Optional[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.
exclusive (bool) – If false (default), all elements are included in the product. If true, the first element is excluded from the product.

Returns:

result – 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.

Return type:

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.script.relax.cumsum(data: Expr, axis: int | None = None, dtype: str | dtype | None = None, exclusive: bool = False, *, ty=None, span=None)#

Numpy style cumsum op. Return the cumulative inclusive sum of the elements along a given axis.

Parameters:

data (relax.Expr) – The input data to the operator.
axis (Optional[int]) – Axis along which the cumulative sum is computed. The default (None) is to compute the cumsum over the flattened array.
dtype (Optional[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.
exclusive (bool) – If false (default), all elements are included in the sum. If true, the first element is excluded from the sum.

Returns:

result – 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.

Return type:

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.script.relax.dataflow(*, span=None)#

Create a dataflow context with explicit finalized exports.

Parameters:: span (SpanEntry, Span or None, optional) – Source location attached to the constructed IR.
Returns:: res – The constructed frame, retaining source metadata.
Return type:: frame.BindingBlockFrame

tvm.script.relax.decl_mutable_cell_(value: ~typing.Any = <MISSING>, *, ty: ~typing.Any | None = None, name: str | None = None, span: ~tvm.script.ir_builder.base.SpanEntry | ~tvm.ir.base.Span | None = None, name_span: ~tvm.script.ir_builder.base.SpanEntry | ~tvm.ir.base.Span | None = None) → NoReturn#

Implements tvm.script.ir_builder.parser_protocol.decl_mutable_cell_().

Relax has immutable bindings and rejects mutable storage declarations.

tvm.script.relax.dequantize(data: Expr, scale: Expr, zero_point: Expr, axis: int = -1, out_dtype: str = 'float32', *, ty=None, span=None)#

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:

data (tvm.relax.Expr) – The input tensor to be dequantized.
scale (tvm.relax.Expr) – The input scale.
zero_point (tvm.relax.Expr) – The input zero_point.
axis (int) – The channel axis for dequantization. Default value is -1 which corresponds to the last axis.
out_dtype (str, optional) – The data type of the output tensor.

Returns:

result – The computed result.

Return type:

tvm.relax.Expr

tvm.script.relax.device(device_type: str | int | DLDeviceType, index: int | None = None) → Device#

Construct a TVM FFI device with given device type and index.

Parameters:

device_type (str or int) – The device type or name.
index (int, optional) – The device index.

Returns:

device

Return type:

tvm_ffi.Device

Examples

Device can be used to create reflection of device by string representation of the device type.

import tvm_ffi

assert tvm_ffi.device("cuda:0") == tvm_ffi.device("cuda", 0)
assert tvm_ffi.device("cpu:0") == tvm_ffi.device("cpu", 0)

tvm.script.relax.device_mesh(shape: Shape, device_ids: list[int] | Range) → DeviceMesh#

Create a device mesh expression. :param shape: The shape of the device mesh. :type shape: Shape :param device_ids: Represents the device id in the mesh :type device_ids: Union[List[int], Range]

Returns:: res – The device mesh.
Return type:: DeviceMesh

tvm.script.relax.divide(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Division with numpy-style broadcasting.

Parameters:

x1 (relax.Expr) – The first input tensor.
x2 (relax.Expr) – The second input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.dtype(value: str | dtype) → DataTypeImm#

Construct a DataType-valued constant for an ordinary operation argument.

Unlike scalar constructors such as T.float32, this constructs a dtype value, rather than a scalar value or variable of that dtype.

tvm.script.relax.dummy_global_info() → DummyGlobalInfo#

Create a dummy global info expression.

Returns:: res – The result dummy global info.
Return type:: DummyGlobalInfo

tvm.script.relax.dynamic_strided_slice(x: Expr, begin: Expr, end: Expr, strides: Expr, *, ty=None, span=None) → Expr#

Dynamic strided slice of a tensor. begin, end, strides can be computed at runtime.

Parameters:

x (Expr) – The source tensor to be sliced.
begin (Expr) – The indices to begin with in the slicing, inclusive.
end (Expr) – The indices indicating end of the slice, exclusive.
strides (Expr) – 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:

ret – The sliced result.

Return type:

Note

dyn_strided_slice require the input begin, end and strides to have the same length as rank of data tensor.

tvm.script.relax.einsum(operands, subscripts, *, ty=None, span=None)#

Evaluates the Einstein summation convention on data

Parameters:

operands (Union(List[relax.Expr], Tuple[relax.Expr])) – A list of expression.
subscripts (str) – The einsum expression string.

Returns:

result – The output from the einsum op.

Return type:

tvm.script.relax.else_(*, span: SpanEntry | Span | None = None) → ElseFrame#: Implements tvm.script.ir_builder.parser_protocol.else_().

tvm.script.relax.emit(value: Expr, annotate_ty: Type | None = None) → Var#

Emit a binding to the last binding block frame. :param value: The right side value of the bindings to be emitted. :type value: Expr :param annotate_ty: The optional type annotation for the emitted value. :type annotate_ty: Optional[Type]

Returns:: var – The left side var of the emitted binding.
Return type:: Var

tvm.script.relax.emit_(value: Any, *, span: SpanEntry | Span | None = None) → None#

Implements tvm.script.ir_builder.parser_protocol.emit_().

Only void expressions may be discarded. Receipts and None emit nothing; non-void expressions raise ValueError and unsupported host values raise TypeError.

tvm.script.relax.emit_match_cast(value: Expr, ty: Type) → Var#

Emit a match_cast binding to the last binding block frame. :param value: The value of the MatchCast to be emitted. :type value: Expr :param ty: The ty of the MatchCast to be emitted. :type ty: Type

Returns:: var – The left side var of the emitted binding.
Return type:: Var

tvm.script.relax.emit_te(func: Callable, *args: Any, **kwargs: Any) → Call#

Emit a call node according to the te function. This function converts arguments from relax expression to te tensor, The callback func should return a te tensor or a list of te tensors.

Parameters:

func (Callable) – A function that returns a te tensor or a list of te tensors.
args (Any, optional) – arguments passed to the function.
kwargs (Any, optional) –
The keyword arguments passed to the function. Note that the following keyword args are reserved:
- ’function_name_hint’ for passing name hint to the Function
  that gets generated.
- ’function_attrs’ is reserved for passing func attributes to
  be added to the Function that gets created.

Returns:

call – A newly created call that calls into a tirx function.

Return type:

tvm.script.relax.emit_var_binding(value: VarBinding) → Var#

Emit a binding to the last binding block frame. :param value: The binding to be emitted. :type value: VarBinding

Returns:: var – The left side var of the emitted binding.
Return type:: Var

tvm.script.relax.emit_with_type(op: str, args: Expr, ty_args: Type | list[Type] | None = None) → Call#

Create a Relax Call with type arguments. :param op: The relax op for which type args are to be appended :type op: Expr :param args: The arguments. :type args: Expr :param ty_args: The list of type arguments. :type ty_args: Optional[Union[Type, List[Type]]]

Returns:: call – The created Relax Call
Return type:: Call

tvm.script.relax.eq_(lhs: Any, rhs: Any, *, span: SpanEntry | Span | None = None) → Expr#: Implements tvm.script.ir_builder.parser_protocol.eq_().

tvm.script.relax.equal(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Broadcasted element-wise test for (lhs == rhs).

Parameters:

x1 (relax.Expr) – The first input tensor.
x2 (relax.Expr) – The second input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.erf(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Computes the error function of the input.

Parameters:: x (relax.Expr) – The input data
Returns:: result – Computed error function for each element.
Return type:: relax.Expr

tvm.script.relax.ewise_fma(x1: Expr, x2: Expr, x3: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Elementwise fused multiply-add operator Returns elementwise result of \(x1 * x2 + x3\)

Parameters:

x1 (relax.Expr) – The left hand operand of the multiplication
x2 (relax.Expr) – The right hand operand of the multiplication
x3 (relax.Expr) – The operand of the addition

Returns:

result – The computed result.

Return type:

tvm.script.relax.exp(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise exp of data.

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

Note

The input tensor is required to have float dtype

tvm.script.relax.expand_dims(x: Expr, axis: int | list[int], *, ty=None, span=None) → Expr#

Insert new axes at the positions given by axis.

Parameters:

x (relax.Expr) – The input data to the operator.
axis (Union[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:

result – The transformed result.

Return type:

tvm.script.relax.ext_dev(dev_id=0)#

Construct a extension device

Parameters:: dev_id (int, optional) – The integer device id
Returns:: dev – The created device
Return type:: Device

Note

This API is reserved for quick testing of new device by plugin device API as ext_dev.

Construct a 2-D tensor with ones on the diagonal and zeros elsewhere.

Parameters:

n (PrimExprLike) – Number of rows in the output.
m (Optional[PrimExprLike]) – Number of columns in the output. If None, defaults to n.
k (PrimExprLike) – 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.
dtype (str | DataType) – The data type of the created tensor.

Returns:

result – The result tensor.

Return type:

tvm.script.relax.eye_like(x: Expr, k: int | Expr = 0, dtype: str | dtype | None = None, *, ty=None, span=None) → Expr#

Return a 2-D tensor with ones on the diagonal and zeros elsewhere, with the same shape as the input tensor.

Parameters:

x (relax.Expr) – The input tensor, which provides the shape, and dtype when the dtype field is not specified.
k (PrimExprLike) – 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.
dtype (Optional[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:

result – The result tensor.

Return type:

tvm.script.relax.flatten(x: Expr, *, ty=None, span=None) → Expr#

Flatten all the tensor dimensions into one.

Parameters:: x (relax.Expr) – The input data to the operator.
Returns:: result – The flattened result.
Return type:: relax.Expr

tvm.script.relax.flip(data, axis, *, ty=None, span=None)#

Reverses the order of elements along given axis while preserving array shape.

Parameters:

data (relax.Expr) – The input data to the operator.
axis (int) – The axis along which to flip over.

Returns:

ret – The computed result.

Return type:

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.script.relax.floor(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Take floor of input data.

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

tvm.script.relax.floor_divide(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Floor division with numpy-style broadcasting.

Parameters:

x1 (relax.Expr) – The first input tensor.
x2 (relax.Expr) – The second input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.floor_mod(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Floor modulo with numpy-style broadcasting.

Parameters:

x1 (Expr) – The first input tensor.
x2 (Expr) – The second input tensor.

tvm.script.relax.for_(iterable: Any, *, names: str | Sequence[str] | None = None, span: SpanEntry | Span | None = None) → NoReturn#

Implements tvm.script.ir_builder.parser_protocol.for_().

Relax rejects source loops; use supported functional control flow.

tvm.script.relax.full(shape: tuple[int | Expr] | Expr, fill_value: Expr, dtype: str | dtype | None = None, *, ty=None, span=None) → Expr#

Fill array with scalar value.

Parameters:

shape (Union[Tuple[PrimExprLike], Expr]) – The shape of the created tensor.
fill_value (relax.Expr) – The value to fill. Must be a scalar tensor.
dtype (Optional[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:

result – The result tensor.

Return type:

tvm.script.relax.full_like(x: Expr, fill_value: Expr, dtype: str | dtype | None = None, *, ty=None, span=None) → Expr#

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:

x (relax.Expr) – The input tensor, which provides the shape, and dtype when the dtype field is not specified.
fill_value (relax.Expr) – The value to fill. Must be a scalar tensor.
dtype (Optional[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:

result – The result tensor.

Return type:

tvm.script.relax.func_attr(attrs: dict[str, Object]) → None#

Specify the attrs of the last function frame.

Parameters:: attrs (Dict[str, Object]) – The function attrs.

tvm.script.relax.func_name_(name: str) → None#: Implements tvm.script.ir_builder.parser_protocol.func_name_().

tvm.script.relax.func_ret_type_(annotation: Any, *, span: SpanEntry | Span | None = None) → None#: Implements tvm.script.ir_builder.parser_protocol.func_ret_type_().

tvm.script.relax.func_ret_value(value: Expr) → None#

Specify the return value of the last function frame.

Parameters:: value (Expr) – The function return value.

tvm.script.relax.function_(pure: bool = True, private: bool = False, *, decl: bool = False, local: bool = False, reference: Var | None = None, span: SpanEntry | Span | None = None) → FunctionFrame#

Implements tvm.script.ir_builder.parser_protocol.function_().

Public pure/private options pass to native purity and visibility controls. Local function bodies require their previously declared reference.

tvm.script.relax.gather_elements(data: Expr, indices: Expr, axis: int = 0, *, ty=None, span=None) → Expr#

Gather elements from data according to indices along the specified axis.

Parameters:

data (relax.Expr) – The input data to the operator.
indices (relax.Expr) – The indices tensor, must have integer type.
axis (int) – The axis along which to index. Default is 0.

Returns:

ret – The computed result.

Return type:

Examples

data = [[1, 2], [3, 4]]
indices = [[0, 0], [1, 0]]
axis = 1
output = [[1, 1], [4, 3]]

data = [[1, 2, 3], [4, 5, 6]]
indices = [[1, 1, 1]]
axis = 0
output = [[4, 5, 6]]

tvm.script.relax.gather_nd(data: Expr, indices: Expr, batch_dims: int = 0, *, ty=None, span=None) → Expr#

Update data at positions defined by indices with values in updates.

Parameters:

data (relax.Expr) – The input data to the operator.
indices (relax.Expr) – The indices tensor, must have integer type.
batch_dims (int) – The number of batch dimensions. Default is 0.

Returns:

ret – The computed result.

Return type:

Examples

batch_dims = 0
data    = [[0,1],[2,3]]   # data_shape    = [2, 2]
indices = [[0,0],[1,1]]   # indices_shape = [2, 2]
output  = [0,3]           # output_shape  = [2]

batch_dims = 1
data    = [[[0,1],[2,3]],[[4,5],[6,7]]] # data_shape    = [2, 2, 2]
indices = [[1],[0]]                     # indices_shape = [2, 1]
output  = [[2,3],[4,5]]                 # output_shape  = [2, 2]

tvm.script.relax.ge_(lhs: Any, rhs: Any, *, span: SpanEntry | Span | None = None) → Expr#: Implements tvm.script.ir_builder.parser_protocol.ge_().

tvm.script.relax.greater(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Broadcasted element-wise test for (lhs > rhs).

Parameters:

x1 (relax.Expr) – The first input tensor.
x2 (relax.Expr) – The second input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.greater_equal(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Broadcasted element-wise test for (lhs >= rhs).

Parameters:

x1 (relax.Expr) – The first input tensor.
x2 (relax.Expr) – The second input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.gt_(lhs: Any, rhs: Any, *, span: SpanEntry | Span | None = None) → Expr#: Implements tvm.script.ir_builder.parser_protocol.gt_().

tvm.script.relax.hamming_window(window_size, periodic, alpha, beta, dtype, *, ty=None, span=None)#

Hamming window function.

Parameters:

window_size (Expr) – The size of returned window.
periodic (Expr) – If True, returns a window to be used as periodic function. If False, return a symmetric window.
alpha (Expr) – The co-efficient alpha.
beta (Expr) – The co-efficient beta.

Returns:

ret – The result tensor.

Return type:

tvm.script.relax.hexagon(dev_id=0)#

Construct a Hexagon device

Parameters:: dev_id (int, optional) – The integer device id
Returns:: dev – The created device
Return type:: Device

tvm.script.relax.hint_on_device(data, device_type, index=0, memory_scope='global', *, ty=None, span=None) → Expr#: Hint the device type, index and memory scope for executing data.

tvm.script.relax.if_(condition: Any, *, span: SpanEntry | Span | None = None) → IfFrame#: Implements tvm.script.ir_builder.parser_protocol.if_().

tvm.script.relax.if_then_else_(condition: Any, true_value: Any, false_value: Any) → Any#: Implements tvm.script.ir_builder.parser_protocol.if_then_else_().

tvm.script.relax.index_put(data: Expr, indices: Expr | tuple[Expr], values: Expr, accumulate: bool = False, *, ty=None, span=None) → Expr#

This operation updates values in data at positions specified by indices with corresponding values from values. The indices is a tuple of tensors where each tensor corresponds to a dimension in data. When accumulate is True, the operation performs accumulation (addition) rather than replacement. The reduction parameter allows specifying different reduction operations. :param data: The input tensor to be modified :type data: relax.Expr :param indices: Tuple of index tensors (one for each dimension) specifying positions to update :type indices: Union[Expr, Tuple[Expr]] :param values: Values to place at the specified indices :type values: relax.Expr :param accumulate: Whether to accumulate (add) values rather than replace (default: False) :type accumulate: bool

Returns:: result – A new tensor with the same shape as data but with specified positions updated
Return type:: relax.Expr

Examples

# inputs
data = torch.zeros(3, 3)
indices = (torch.tensor([0, 2]), torch.tensor([1, 1]))
values = torch.tensor([1.0, 2.0])
# output
output = [
    [0.0, 1.0, 0.0],
    [0.0, 0.0, 0.0],
    [0.0, 2.0, 0.0],
]
# with accumulate=True
output = [
    [0.0, 1.0, 0.0],
    [0.0, 0.0, 0.0],
    [0.0, 3.0, 0.0],
]

tvm.script.relax.index_tensor(data: Expr, indices: Expr | list[Expr], *, ty=None, span=None) → Expr#

Advanced-tensor indexing (NumPy/PyTorch-style).

Given k index tensors indices = (I0, I1, …, Ik-1) this operator selects elements from data as if one had written data[I0, I1, …, Ik-1] in NumPy/PyTorch:

All index tensors must have an integer dtype.

Their shapes are broadcast together to a common shape B in the usual NumPy way.

The result shape is B + data.shape[k:] (i.e. the broadcast shape followed by the remaining axes of data that are not indexed).

At compile-time Relax checks that the number of index tensors k does not exceed data.ndim, that the dtypes are integer, and that the shapes are consitent (broadcast-compatible).

Parameters:

data (relax.Expr) – The input tensor to be indexed.
indices (Union[relax.Expr, List[relax.Expr]]) – A Tuple expression containing the index tensors, or a Python list / tuple that will be promoted to a tuple expression automatically. Each tensor must have an integer dtype.

Returns:

result – The tensor obtained after advanced indexing. Its dtype equals data.dtype

Return type:

Examples

import numpy as np
import tvm.relax as R

x   = R.const(np.arange(9).reshape(3, 3).astype("float32"))
row = R.const(np.array([0, 2]))        # shape (2,)
col = R.const(np.array([1, 0]))        # shape (2,)

y = R.index_tensor(x, [row, col])
# y.shape == (2,) ;  y == [1., 6.]

# Broadcasting: row : (2,1), col : (1,3)  →  B = (2,3)
row = R.const(np.array([[0],[1]]))
col = R.const(np.array([[0,1,2]]))
z = R.index_tensor(x, [row, col])
# z.shape == (2,3)

tvm.script.relax.invoke_closure(closure: Expr, args: Expr, ty_args: list[Type] | Type, *, ty=None, span=None) → Call#

Invoke a closure.

Parameters:

closure (Expr) – The VMClosure object.
args (Expr) – The input arguments.
type_args (Union[List[Type], Type]) – The type information arguments of the CallNode

Returns:

ret – A call to invoke_closure.

Return type:

tvm.script.relax.invoke_pure_closure(closure: Expr, args: Expr, ty_args: list[Type] | Type, *, ty=None, span=None) → Call#

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:

closure (Expr) – The VMClosure object.
args (Expr) – The input arguments.
type_args (Union[List[Type], Type]) – The type information arguments of the CallNode

Returns:

ret – A call to invoke_pure_closure.

Return type:

tvm.script.relax.isfinite(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Check if input value is finite.

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

tvm.script.relax.isinf(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Check if input value is infinite.

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

tvm.script.relax.isnan(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Check if input value is Nan.

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

tvm.script.relax.layout_transform(x: Expr, index_map: Callable | IndexMap, pad_value: int | float | Expr | None = None, *, ty=None, span=None)#

Modifies the layout of a tensor.

Parameters:

x (relax.Expr) – The input tensor to the operator.
index_map (Callable | IndexMap) – The transformation to apply.
pad_value (Optional[int | float | Expr]) – The value used for padding if the transformation results in implicit padding. If not specified, any value can be used.

Returns:

result – The transformed tensor.

Return type:

tvm.script.relax.le_(lhs: Any, rhs: Any, *, span: SpanEntry | Span | None = None) → Expr#: Implements tvm.script.ir_builder.parser_protocol.le_().

tvm.script.relax.left_shift(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Bitwise Shift Left :param x1: The input tensor to be shifted. :type x1: relax.Expr :param x2: The number of positions to shift. :type x2: relax.Expr

Returns:: result – The computed result.
Return type:: relax.Expr

tvm.script.relax.less(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Broadcasted element-wise test for (lhs < rhs).

Parameters:

x1 (relax.Expr) – The first input tensor.
x2 (relax.Expr) – The second input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.less_equal(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Broadcasted element-wise test for (lhs <= rhs).

Parameters:

x1 (relax.Expr) – The first input tensor.
x2 (relax.Expr) – The second input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.linear(data: Expr, weight: Expr, bias: Expr | None = None, out_dtype: str | dtype | None = None) → Expr#

Applies a linear transformation to the incoming data: y = xA^T + b

Parameters:

data (relax.Expr) – The input data.
weight (relax.Expr) – The weight tensor.
bias (Optional[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:: result – The computed result.
Return type:: relax.Expr

tvm.script.relax.log(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise natural logarithm of the input data.

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

Note

The input tensor is required to have float dtype

tvm.script.relax.log_add_exp(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute the log of the sum of exponentials of the inputs, element-wise.

Parameters:

x1 (Expr) – The first input tensor.
x2 (Expr) – The second input tensor.

Returns:

The element-wise log-sum-exp of x1 and x2.

Return type:

tvm.script.relax.logical_and(*values, ty_args=None, ty=None, span=None)#

Construct conjunction of host, primitive, or tensor values.

Parameters:: values (Expr or Python value) – One or more eagerly evaluated operands. Host pairs follow Python logical operations; primitive pairs use primitive IR and tensor pairs use the corresponding Relax operation.
Returns:: result – The conjunction reduced from left to right.
Return type:: Expr or Python value

Notes

All arguments are evaluated before this call; it does not provide Python short-circuit evaluation of the argument expressions.

tvm.script.relax.logical_not(value, *, ty_args=None, ty=None, span=None)#

Negate a host, primitive, or tensor condition.

Parameters:: value (Expr or Python value) – Operand to negate. Primitive and tensor expressions use their IR logical operation; host values use Python truth testing.
Returns:: result – The logical negation without testing an IR expression as a Python bool.
Return type:: Expr or bool

tvm.script.relax.logical_or(*values, ty_args=None, ty=None, span=None)#

Construct disjunction of host, primitive, or tensor values.

Parameters:: values (Expr or Python value) – One or more eagerly evaluated operands. Host pairs follow Python logical operations; primitive pairs use primitive IR and tensor pairs use the corresponding Relax operation.
Returns:: result – The disjunction reduced from left to right.
Return type:: Expr or Python value

Notes

All arguments are evaluated before this call; it does not provide Python short-circuit evaluation of the argument expressions.

tvm.script.relax.logical_xor(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Logical XOR :param x1: The first input tensor. :type x1: relax.Expr :param x2: The second input tensor. :type x2: relax.Expr

Returns:: result – The computed result.
Return type:: relax.Expr

tvm.script.relax.lookup_vdevice(target_kind: str | None = None, device_index: int = -1) → VDevice#

Retrieve a virtual device from the active module’s global-info list.

Parameters:

target_kind (str) – The target device kind, for example ‘llvm’ or ‘cuda’. Use ‘vdevice’ to index the complete virtual-device list.
device_index (int) – The zero-based index among devices of the selected target kind, or among all devices when target_kind is ‘vdevice’.

Returns:

res – The result virtual device.

Return type:

VDevice

tvm.script.relax.lt_(lhs: Any, rhs: Any, *, span: SpanEntry | Span | None = None) → Expr#: Implements tvm.script.ir_builder.parser_protocol.lt_().

tvm.script.relax.make_closure(func: Expr, args: Expr, *, ty=None, span=None) → Object#

Create a closure with free variables and return the closure.

Parameters:

func (Expr) – The closure, can be ExternFunc or Function.
args (Expr) – The input arguments.

Returns:

ret – The VMClosure.

Return type:

Object

tvm.script.relax.match_cast(value, ty, *, span=None)#

Construct a match-cast descriptor for the binding hook.

Parameters:

value (Expr or Python value) – Value to match against the asserted type. Numbers and Python tuples are converted to Relax expressions; None is not accepted.
ty (Type or callable) – Asserted type, or a zero-argument factory producing its annotation.
span (SpanEntry, Span or None, optional) – Source location attached to the constructed IR; None leaves it unspecified.

Returns:

result – An unbound match-cast descriptor, consumed by the language variant binding hook to emit and name the result.

Return type:

MatchCast

tvm.script.relax.matmul(x1: Expr, x2: Expr, out_dtype: str | dtype | None = None, *, ty=None, span=None) → Expr#

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:

x1 (relax.Expr) – The first input tensor.
x2 (relax.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:

result – The computed result.

Return type:

tvm.script.relax.max(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False, *, ty=None, span=None) → Expr#

Computes the max of tensor elements over given axes.

Parameters:

x (relax.Expr) – The input data tensor
axis (Optional[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.
keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.maximum(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Element-wise maximum

Parameters:

x1 (relax.Expr) – The first input tensor.
x2 (relax.Expr) – The second input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.mean(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False, *, ty=None, span=None) → Expr#

Computes the mean of tensor elements over given axes.

Parameters:

x (relax.Expr) – The input data tensor
axis (Optional[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.
keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.median(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False, *, ty=None, span=None) → Expr#

Computes the median of tensor elements over given axes.

Parameters:

x (relax.Expr) – The input data tensor
axis (Optional[Union[int, List[int]]]) – Axis along which the median is computed. The default (None) is to compute the median of the entire flattened tensor.
keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.meshgrid(tensors: Expr | list[Expr], indexing: str | None = 'ij', *, ty=None, span=None) → Expr#

Generate coordinate grids from input tensors.

Parameters:

tensors (Union[relax.Expr, List[relax.Expr]]) – An Expr in Tuple type, containing 1D tensors (or scalars promoted to 1D) to generate coordinate grids from, or a list of such tensors.
indexing (Optional[str]) – The indexing mode, either “ij” (matrix indexing) or “xy” (Cartesian indexing). Defaults to “ij”.

Returns:

result – A Tuple of tensors representing the coordinate grids.

Return type:

tvm.script.relax.metal(dev_id=0)#

Construct a metal device

Parameters:: dev_id (int, optional) – The integer device id
Returns:: dev – The created device
Return type:: Device

tvm.script.relax.min(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False, *, ty=None, span=None) → Expr#

Computes the min of tensor elements over given axes.

Parameters:

x (relax.Expr) – The input data tensor
axis (Optional[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.
keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.minimum(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Element-wise minimum

Parameters:

x1 (relax.Expr) – The first input tensor.
x2 (relax.Expr) – The second input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.mod(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Modulo with numpy-style broadcasting.

Parameters:

x1 (Expr) – The first input tensor.
x2 (Expr) – The second input tensor.

tvm.script.relax.multinomial_from_uniform(prob: Expr, uniform_sample: Expr, sample_indices: Expr, dtype: str = 'int64', *, ty=None, span=None) → Expr#

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:

prob (relax.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_sample (relax.Expr) – The uniformly sampled 2-D tensor with the shape (n, 1). Values range from 0 to 1, indicating probabilities sampled uniformly.
sample_indices (relax.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].
dtype (str) – The data type of the output tensor.

Returns:

result – The computed tensor with shape (n, 1).

Return type:

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]]

tvm.script.relax.multiply(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Multiplication with numpy-style broadcasting.

Parameters:

x1 (Expr) – The first input tensor.
x2 (Expr) – The second input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.ne_(lhs: Any, rhs: Any, *, span: SpanEntry | Span | None = None) → Expr#: Implements tvm.script.ir_builder.parser_protocol.ne_().

tvm.script.relax.negative(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise negative of the input data.

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

tvm.script.relax.nonzero(x: Expr, *, ty=None, span=None) → Expr#

Find the indices of elements of a tensor that are non-zero.

Parameters:: x (relax.Expr) – The input data tensor.
Returns:: result – A 2-D tensor containing indices of non-zero elements.
Return type:: relax.Expr

Note

This function is equivalent to onnx.nonzero.

Examples

x = [[0, 1],
     [2, 0]]
nonzero(x) = [[0, 1],
              [1, 0]]

tvm.script.relax.not_(value: Any) → Any#: Implements tvm.script.ir_builder.parser_protocol.not_().

tvm.script.relax.not_equal(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Broadcasted element-wise test for (lhs != rhs).

Parameters:

x1 (relax.Expr) – The first input tensor.
x2 (relax.Expr) – The second input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.null_value(*, ty=None, span=None) → Call#

Create a call node that represents a null value object.

Returns:: ret – The created call node.
Return type:: Call

tvm.script.relax.one_hot(indices: Expr, on_value: int | float | Expr, off_value: int | float | Expr, depth: int, axis: int = -1, *, ty=None, span=None) → Expr#

Returns a one-hot tensor.

Parameters:

indices (relax.Expr) – The indices to set to on_value.
on_value (int | float | Expr) – The value to fill at indices.
off_value (int | float | Expr) – The value to fill at other locations.
depth (int) – The depth of the one-hot dimension.
axis (int, optional) – The axis to fill. Default is -1 which adds a new dimension at the end.

Returns:

result – The computed result.

Return type:

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.script.relax.ones(shape: tuple[int | Expr] | Expr, dtype: str | dtype, *, ty=None, span=None) → Expr#

Construct a tensor of all ones, with the input shape and dtype.

Parameters:

shape (Union[Tuple[PrimExprLike], Expr]) – The shape of the created tensor.
dtype (str | DataType) – The data type of the created tensor.

Returns:

result – The result tensor.

Return type:

tvm.script.relax.ones_like(x: Expr, dtype: str | dtype | None = None, *, ty=None, span=None) → Expr#

Construct a tensor with all ones, with shape of the input tensor shape.

Parameters:

x (relax.Expr) – The input tensor, which provides the shape, and dtype when the dtype field is not specified.
dtype (Optional[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:

result – The result tensor.

Return type:

tvm.script.relax.opencl(dev_id=0)#

Construct a OpenCL device

Parameters:: dev_id (int, optional) – The integer device id
Returns:: dev – The created device
Return type:: Device

tvm.script.relax.or_(*values: Any) → Any#: Implements tvm.script.ir_builder.parser_protocol.or_().

tvm.script.relax.outer(x1: Expr, x2: Expr, *, ty=None, span=None) → Expr#

Computes the outer product of two input expressions.

Parameters:

x1 (relax.Expr) – The first input expression.
x2 (relax.Expr) – The second input expression.

Notes

This operation computes the outer product between two expressions, resulting in a tensor where each element is the product of elements from x1 and x2. It is commonly used in tensor and matrix operations to expand lower-dimensional inputs into higher-dimensional representations.

Returns:: result – The resulting expression representing the outer product.
Return type:: relax.Expr

tvm.script.relax.output(*vars: tuple[Var]) → None#

Expose the dataflow block output variables as global ones.

Parameters:: vars (Tuple[Var]) – The output variables of a dataflow block.

tvm.script.relax.permute_dims(x: Expr, axes: list[int] | None = None, *, ty=None, span=None) → Expr#

Permutes the dimensions of an array.

Parameters:

x (relax.Expr) – The input data to the operator.
axes (Optional[List[int]]) – The target axes order. If not specified, permute_dims will reverse the order of all axes.

Returns:

result – The transposed result.

Return type:

tvm.script.relax.power(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None)#

Power with numpy-style broadcasting.

Parameters:

x1 (relax.Expr) – The first input tensor.
x2 (relax.Expr) – The second input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.prim_value(value: Expr | int | float, dtype: str | None = None) → Expr#

Convert a Python scalar or primitive expression to Expr.

Parameters:

value (Expr | int | float) – The value to convert.
dtype (Optional[str]) – The dtype to use when converting Python numeric values.

Returns:

result – The converted primitive expression. Existing primitive expression inputs are returned unchanged, even when dtype is supplied.

Return type:

Notes

Existing primitive expressions can be used directly in TVMScript. Python scalars are context dependent: an unannotated Relax binding or return converts them to tensor constants, while this helper constructs primitive expressions. dtype selects the type of Python numeric inputs; omitting it uses bool, int64, or float64. The script alias R.prim_value retains these programmatic conversion semantics.

tvm.script.relax.print(format: str | Expr, *values: Expr, ty=None, span=None) → Expr#: Print values using the canonical leading format-string operand.

tvm.script.relax.prod(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False, *, ty=None, span=None) → Expr#

Computes the product of tensor elements over given axes.

Parameters:

x (relax.Expr) – The input data tensor
axis (Optional[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.
keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.py_module(module: type | None = None, **options: _Any) → IRModule | _PyModuleFactory | _Callable[[type], IRModule | _PyModuleFactory]#

Parse a class and attach executable Relax Python-function metadata.

Parameters:

module (type, optional) – Class containing registered script functions and I.pyfunc members. None, the default, returns a decorator. Python bodies are retained as their original callables and are not executed during construction.
**options (_Any) – Options forwarded once to shared parse, including check_well_formed (True by default) and track_span (True).

Returns:

An IRModule with original Python callables in __pyfuncs__ and corresponding Relax ExternFunc metadata. For a BasePyModule subclass, return a factory accepting optional device and target, creating a fresh runtime instance on each call. With no class, return the decorator.

Return type:

IRModule or callable

Raises:

TypeError – If the decorated value is not a class, or an unsupported parser option is supplied.
SyntaxError – If the source violates a parser syntax restriction. Construction and enabled validation propagate their original exceptions unchanged.
OSError – If source inspection cannot recover a Python function’s definition.

Notes

The decorator captures the original class-definition scope and passes it explicitly to shared parsing. It retains no Python frame or scope snapshot. ExternFunc metadata uses each original function’s source coordinates when span tracking is enabled. Runtime compilation/registration is deferred to BasePyModule construction; plain IRModule construction registers no runtime functions. Ordinary shared I.ir_module never performs these Relax steps.

# Source
@R.py_module
class Module:
    @I.pyfunc
    def twice(value):
        return value * 2

# Relax adaptation after one shared parse
assert Module.__pyfuncs__["twice"](3) == 6
assert isinstance(Module["twice"], relax.ExternFunc)

tvm.script.relax.quantize(data: Expr, scale: Expr, zero_point: Expr, axis: int = -1, out_dtype: str = 'int8', *, ty=None, span=None)#

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:

data (tvm.relax.Expr) – The input tensor to be quantized.
scale (tvm.relax.Expr) – The output scale.
zero_point (tvm.relax.Expr) – The output zero_point.
axis (int) – The channel axis for quantization. Default value is -1 which corresponds to the last axis.
out_dtype (str, optional) – The data type of the output tensor.

Returns:

result – The computed result.

Return type:

tvm.relax.Expr

tvm.script.relax.range_(*args: Any, annotations: dict[str, Any] | None = None) → NoReturn#: Implements tvm.script.ir_builder.parser_protocol.range_().

tvm.script.relax.repeat(data: Expr, repeats: int, axis: int | None = None, *, ty=None, span=None) → Expr#

Repeats elements of an array.

Parameters:

data (relax.Expr) – The input tensor.
repeats (int) – 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:

ret – The computed result.

Return type:

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.script.relax.reshape(x: Expr, shape: tuple[int | Expr] | Expr, *, ty=None, span=None) → Expr#

Reshape the input array.

-1 infers the dimension of the output shape by using the remainder of the input dimensions keeping the size of the new array same as that of the input array. At most one dimension of shape can be -1.

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:

x (relax.Expr) – The input data to the operator.
shape (Union[Tuple[PrimExprLike], Expr]) – The new shape. Should be compatible with the original shape.

Returns:

result – The reshaped result.

Return type:

Note

The -1 inference is only performed at compile-time. That is to say, in any case the dimension length of -1 cannot be inferred in compile-time, an error will be thrown.

tvm.script.relax.resolve_global_info_(content: str) → GlobalInfo#

Resolve a module-owned named-list or virtual-device selector.

Parameters:: content (str) – Original global-info selector. “mesh[0]” indexes a named list; “cuda:1” selects the second CUDA vdevice; “vdevice:0” selects by absolute index. A trailing memory-scope suffix is accepted without changing device selection.
Returns:: The exact registered global-info object.
Return type:: GlobalInfo

Notes

Lookup requires the nearest active native module frame and creates no metadata. Missing context, malformed selectors or unmatched devices raise ValueError; missing map entries or out-of-range indices propagate KeyError/IndexError. Non-string inputs raise TypeError; constructor argument handling preserves concrete objects before calling this hook. No source span is attached to an existing metadata object.

# The constructor decorator calls this resolver for string selectors.
R.Tensor((n,), "float32", vdevice="cuda:0")
# Direct resolution uses the same module metadata.
device = R.resolve_global_info_("cuda:0")

tvm.script.relax.return_(value: Any | None = None, *, span: SpanEntry | Span | None = None) → None#: Implements tvm.script.ir_builder.parser_protocol.return_().

tvm.script.relax.reverse_sequence(data: Expr, seq_lengths: Expr, seq_axis: int = 1, batch_axis: int = 0, *, ty=None, span=None) → Expr#

Reverses variable length slices.

Parameters:

data (relax.Expr) – The input tensor.
seq_lengths (relax.Expr) – A 1-D tensor containing sequence lengths for each batch.
seq_axis (int) – The axis along which to reverse variable length slices.
batch_axis (int) – The axis that indexes the batch.

Returns:

ret – The computed result.

Return type:

tvm.script.relax.rewriter(rewriter_mod: IRModule | type) → PatternMatchingRewriter#

Define a pattern-rewrite rule

The IRModule must have two publicly-exposed functions, pattern and replacement, where pattern and replacement have the same function signature.

@R.rewriter
class RewriteAddIntoMultiply:
    @R.function
    def pattern(A: R.Tensor):
        B = A + A
        return B

    @R.function
    def replacement(A: R.Tensor):
        B = A * 2
        return B

Parameters:: rewriter_mod (Union[IRModule, Type]) – Either an IRModule that defines a rewrite pattern, or a TVMScript class that can be parsed into an IRModule.
Returns:: rewriter – A rewriter object, which can be applied either to a Relax function or to an entire IRModule.
Return type:: PatternMatchingRewriter

Notes

Class members are parsed together after the class body completes. Their annotations use the decorator’s original definition scope, which is released after parsing. An existing IRModule is used directly.

tvm.script.relax.right_shift(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Bitwise Shift Right :param x1: The input tensor to be shifted. :type x1: relax.Expr :param x2: The number of positions to shift. :type x2: relax.Expr

Returns:: result – The computed result.
Return type:: relax.Expr

tvm.script.relax.rocm(dev_id=0)#

Construct a ROCM device

Parameters:: dev_id (int, optional) – The integer device id
Returns:: dev – The created device
Return type:: Device

tvm.script.relax.round(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Rounds each element of the input data to nearest integer.

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

tvm.script.relax.rsqrt(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise reciprocal square root of the input data.

\[1/sqrt(x)\]

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

Note

The input tensor is required to have float dtype

tvm.script.relax.scatter_elements(data: Expr, indices: Expr, updates: Expr, axis: int = 0, reduction: str = 'update', *, ty=None, span=None)#

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:

data (relax.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:

result – The result has the same size as data, and the same shape as data

Return type:

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.script.relax.scatter_nd(data: Expr, indices: Expr, updates: Expr, reduction: str = 'update', *, ty=None, span=None) → Expr#

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:

result – The result has the same shape as data.

Return type:

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]

tvm.script.relax.select(condition, true_value, false_value)#

Select between already-constructed values.

Parameters:

condition (Expr or Python value) – Primitive condition, tensor condition, or host truth value.
true_value (Expr or Python value) – Value selected when the condition is true.
false_value (Expr or Python value) – Value selected when the condition is false.

Returns:

result – A primitive Select, Relax elementwise where, or the selected host value.

Return type:

Expr or Python value

Notes

Both value arguments are evaluated before this call. Host selection returns the selected object; tensor selection converts Python numbers and tuples.

tvm.script.relax.seq_expr() → SeqExprFrame#

Create a SeqExpr frame.

Returns:: res – The result SeqExprFrame
Return type:: _frame.SeqExprFrame

tvm.script.relax.set_mutable_cell_(target: Any, value: Any, *, span: SpanEntry | Span | None = None) → NoReturn#

Implements tvm.script.ir_builder.parser_protocol.set_mutable_cell_().

Relax has immutable bindings and rejects mutable storage updates.

tvm.script.relax.setattr_(target: Any, name: str, value: Any, *, span: SpanEntry | Span | None = None) → NoReturn#: Implements tvm.script.ir_builder.parser_protocol.setattr_().

tvm.script.relax.setitem_(target: Any, key: Any, value: Any, *, span: SpanEntry | Span | None = None) → NoReturn#: Implements tvm.script.ir_builder.parser_protocol.setitem_().

tvm.script.relax.shape#: alias of ShapeExpr

tvm.script.relax.shape_of(expr: Expr, *, ty=None, span=None) → Expr#

Get shape of a tensor.

Parameters:: expr (Expr) – The input Expr.
Returns:: result – A relax Call, which gets the shape of the input
Return type:: Expr

tvm.script.relax.shape_to_tensor(expr: Expr, *, ty=None, span=None) → Expr#

Convert shape to tensor expr. :param expr: The input Expr :type expr: Expr

Returns:: result – A relax Call, which transforms the shape values to the tensor
Return type:: Expr

tvm.script.relax.sigmoid(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise sigmoid of the input data.

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

Note

The input tensor is required to have float dtype

tvm.script.relax.sign(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Returns an indication of the sign of a number for each element of the input data.

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

tvm.script.relax.sin(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise sin of the input data.

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

Note

The input tensor is required to have float dtype

tvm.script.relax.sinh(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise sinh of the input data.

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

Note

The input tensor is required to have float dtype

tvm.script.relax.size(expr: Expr, *, ty=None, span=None) → Expr#

Get the total number of elements in a tensor.

Parameters:: expr (Expr) – The input tensor.
Returns:: result – A scalar tensor of dtype int64 containing the total number of elements.
Return type:: Expr

tvm.script.relax.slice_scatter(input_tensor: Expr, src: Expr, start, end, step, axis=0, *, ty=None, span=None)#

Embeds the values of the src tensor into input at the given dimension.

Parameters:

input_tensor (relax.Expr) – The input tensor to be updated.
src (relax.Expr) – The tensor to embed into input.
axis (int) – The dimension to insert the slice into.
start – The start index of where to insert the slice.
end – The end index of where to insert the slice.
step – The how many elements to skip in.

Returns:

result – The computed result tensor with the same shape as data.

Return type:

tvm.script.relax.sort(x: Expr, axis: int = -1, descending: bool = False, *, ty=None, span=None)#

Performs sorting along the given axis and returns an array in sorted order.

Parameters:

x (relax.Expr) – The input tensor.
axis (int) – Axis along which to sort the input tensor. By default the last axis of the input is used.
descending (bool) – Whether to sort in descending order, the default is False

Returns:

out – Sorted tensor.

Return type:

tvm.script.relax.split(x: Expr, indices_or_sections: int | list[int | Expr], axis: int = 0, *, ty=None, span=None) → Expr#

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 Expr, the entries indicate the indices where along axis the array is split.

Parameters:

x (relax.Expr) – The tensor to be split.
indices_or_sections (Union[int, List[PrimExprLike]]) – Indices or sections to split into. Accepts an int or a list.
axis (int) – The axis over which to split.

Returns:

ret – The computed result.

Return type:

tvm.script.relax.sqrt(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise square root of the input data.

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

Note

The input tensor is required to have float dtype

tvm.script.relax.square(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Squares each element of the input data.

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

tvm.script.relax.squeeze(x: Expr, axis: int | list[int] | None = None, *, ty=None, span=None) → Expr#

Squeeze axes in the array.

Parameters:

x (relax.Expr) – The input data to the operator.
axis (Optional[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:

result – The squeezed result.

Return type:

tvm.script.relax.stack(tensors: Expr | list[Expr], axis: int = 0, *, ty=None, span=None) → Expr#

Stack the input tensors along a new axis.

Parameters:

tensors (Union[relax.Expr, List[relax.Expr]]) – An Expr in Tuple type, containing the tensors to be stacked, or a list of Tensors. All input tensors must have the same shape.
axis (int) – The axis in the resulting tensor along which the input tensors will be stacked. Negative values wrap around. Default is 0.

Returns:

result – The stacked tensor with an additional dimension compared to the input tensors.

Return type:

tvm.script.relax.std(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False, *, ty=None, span=None) → Expr#

Computes the standard deviation of tensor elements over given axes.

Parameters:

x (relax.Expr) – The input data tensor
axis (Optional[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.
keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.stop_lift_params(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

An indicator that the consumers of input tensor should not be lifted to transform_params function

Parameters:: x (relax.Expr) – The input data
Returns:: result – The result tensor that is the same as input tensor
Return type:: relax.Expr

tvm.script.relax.str#: alias of StringImm

tvm.script.relax.strided_slice(x: Expr, axes: Expr, begin: Expr, end: Expr, strides: Expr | None = None, assume_inbound: bool = False, *, ty=None, span=None) → Expr#

Strided slice of a tensor.

Parameters:

x (relax.Expr) – The source tensor to be sliced.
axes (List[int]) – Axes along which slicing is applied.
begin (List[PrimExprLike]) – The indices to begin with in the slicing, inclusive.
end (List[PrimExprLike]) – The indices indicating end of the slice, exclusive.
strides (Optional[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_inbound (bool) – 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:

ret – The sliced result.

Return type:

Note

strided_slice require the input begin, end and strides to have the same length as axes.

tvm.script.relax.subtract(x1: Expr, x2: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Subtraction with numpy-style broadcasting.

Parameters:

x1 (relax.Expr) – The first input tensor.
x2 (relax.Expr) – The second input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.sum(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False, *, ty=None, span=None) → Expr#

Computes the sum of tensor elements over given axes.

Parameters:

x (relax.Expr) – The input data tensor
axis (Optional[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.
keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.take(x: Expr, indices: Expr, axis: int | None = None, mode: str = 'fast', *, ty=None, span=None) → Expr#

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:

x (relax.Expr) – The source tensor.
indices (relax.Expr) – The indices of the values to extract.
axis (Optional[int]) – The axis over which to select values. If it is none, the input tensor is required to be one-dimensional.
mode (str) – Specifies how out-of-bounds indices will behave. - fast (default): extra indices lead to seg fault (user must make sure indices are in-bound) - nan: produce NaNs for out-of-bounds indices - wrap: wrap around the indices - clip: clip to the range

Returns:

ret – The taken result.

Return type:

tvm.script.relax.tan(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise tan of the input data.

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

Note

The input tensor is required to have float dtype

tvm.script.relax.tanh(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Compute element-wise tanh of the input data.

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

Note

The input tensor is required to have float dtype

tvm.script.relax.tensor_to_shape(expr: Expr, *, ty=None, span=None) → Expr#

Convert tensor to shape expr. :param expr: The input Expr :type expr: Expr

Returns:: result – A relax Call, which transforms the tensor values to the shape
Return type:: Expr

tvm.script.relax.then_(*, span: SpanEntry | Span | None = None) → ThenFrame#: Implements tvm.script.ir_builder.parser_protocol.then_().

tvm.script.relax.tile(data: Expr, repeats: int | tuple[int] | list[int], *, ty=None, span=None) → Expr#

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:

data (relax.Expr) – The input data to the operator.
repeats (Union[int, Tuple[int], List[int]]) – The number of repetitions of data along each axis.

Returns:

ret – The computed result.

Return type:

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.script.relax.to_vdevice(data: Expr, dst_vdevice: str | VDevice | None = None, *, ty=None, span=None) → Expr#

Copy data to the destination device.

Parameters:

data (Expr) – The tensor to be copied.
dst_vdevice (Union[py_str, VDevice]) – The destination device where the data is copied to.

Returns:

result – The copied result.

Return type:

tvm.script.relax.topk(data: Expr, k: int = 1, axis: int = -1, ret_type: str = 'both', largest: bool = True, dtype: str = 'int32', *, ty=None, span=None)#

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:

data (relax.Expr) – The input data tensor.
k (int) – Number of top elements to select. Return all elements if k < 1.
axis (int) – 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.
largest (bool) – Whether to return largest or smallest elements. The k smallest elements are returned if largest is False.
dtype (str) – The data type of the indices output.

Returns:

out – The computed result.

Return type:

relax.Expr or List[relax.Expr]

tvm.script.relax.tril(x: Expr, k: int | Expr = 0, *, ty=None, span=None) → Expr#

Return the lower triangular part of a matrix or a batch of matrices.

Parameters:

x (relax.Expr) – The tensor that tril will be applied to. It is required to have at least two dimensions.
k (int) – 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:

ret – The result tensor.

Return type:

tvm.script.relax.triu(x: Expr, k: int | Expr = 0, *, ty=None, span=None) → Expr#

Return the upper triangular part of a matrix or a batch of matrices.

Parameters:

x (relax.Expr) – The tensor that triu will be applied to. It is required to have at least two dimensions.
k (int) – 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:

ret – The result tensor.

Return type:

tvm.script.relax.trunc(x: Expr, *, ty_args=None, ty=None, span=None) → Expr#

Take trunc of input data. :param x: The input data :type x: relax.Expr

Returns:: result – The computed result.
Return type:: relax.Expr

tvm.script.relax.tuple(*fields: Expr) → Expr#

Create a tuple expression from already constructed expression fields.

In script bindings and returns, prefer ordinary Python tuple syntax, including () and (value,), when fields are already expressions. This helper materializes an IR tuple outside parser conversion contexts. Numeric fields use core constructor conversion: for example, tuple(1) has an int32 field, whereas a script tuple literal (1,) has an int64 field. Use explicitly typed expressions to make either spelling equivalent.

Parameters:: *fields (Expr) – The fields of the tuple.
Returns:: res – The result tuple.
Return type:: Expr

tvm.script.relax.type_var(name, *, dtype=None, span=None)#

Construct a fresh standalone primitive symbol.

Parameters:

name (str) – Name of the symbol.
dtype (str or PrimType, optional) – Primitive type of the symbol; None selects “int64”.
span (SpanEntry, Span or None, optional) – Source location attached to the constructed IR; None leaves it unspecified.

Returns:

result – The newly constructed primitive variable.

Return type:

Notes

This constructor creates a new symbol on each call. Reuse the returned variable to share its identity across annotations and function parameters.

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:

x (relax.Expr) – The input tensor.
sorted (Union[bool, Expr]) – Whether to sort the unique elements in ascending order before returning as output.
return_index (Union[bool, Expr]) – Whether to return an additional tensor with indices for where elements in the unique tensor come from the original input.
return_inverse (Union[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_counts (Union[bool, Expr]) – Whether to return an additional tensor with counts of each unique elements.
axis (Optional) – The dimension to apply unique. If not specified, the unique values of the flattened input are returned.

Returns:

ret – The created relax call with

Return type:

tvm.script.relax.unpack_(value: Any) → Any#: Implements tvm.script.ir_builder.parser_protocol.unpack_().

tvm.script.relax.variance(x: Expr, axis: int | list[int] | None = None, keepdims: bool = False, *, ty=None, span=None) → Expr#

Computes the variance of tensor elements over given axes.

Parameters:

x (relax.Expr) – The input data tensor
axis (Optional[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.
keepdims (bool) – If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input tensor.

Returns:

result – The computed result.

Return type:

tvm.script.relax.vdevice(target=None, vdevice_id: int = 0, memory_scope: str = 'global') → VDevice#

Create a virtual device global info. :param target: The target. :param vdevice_id: The virtual device index. :type vdevice_id: int :param memory_scope: The memory scope, default is “global” :type memory_scope: py_str

Returns:: res – The result virtual device.
Return type:: VDevice

tvm.script.relax.vpi(dev_id=0)#

Construct a VPI simulated device

Parameters:: dev_id (int, optional) – The integer device id
Returns:: dev – The created device
Return type:: Device

tvm.script.relax.vulkan(dev_id=0)#

Construct a Vulkan device

Parameters:: dev_id (int, optional) – The integer device id
Returns:: dev – The created device
Return type:: Device

tvm.script.relax.webgpu(dev_id=0)#

Construct a webgpu device.

Parameters:: dev_id (int, optional) – The integer device id
Returns:: dev – The created device
Return type:: Device

tvm.script.relax.where(condition: Expr, x1: Expr, x2: Expr, *, ty=None, span=None) → Expr#

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:

condition (relax.Expr) – When True, yield x1; otherwise, yield x2. Must be broadcasting compatible with x1 and x2. Must have boolean dtype.
x1 (relax.Expr) – The first input tensor. Must be broadcasting compatible with condition and x2.
x2 (relax.Expr) – The second input tensor. Must be broadcasting compatible with condition and x1.

Returns:

result – The result tensor.

Return type:

tvm.script.relax.while_(condition: Any, *, span: SpanEntry | Span | None = None) → NoReturn#: Implements tvm.script.ir_builder.parser_protocol.while_().

tvm.script.relax.wrap_param(data: Expr, dtype: str | dtype | PrimType = 'float32', *, ty=None, span=None) → Expr#

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. :param data: The input data to the operator. :type data: relax.Expr :param dtype: The target data type :type dtype: Union[str, DataType]

Returns:: result – The casted result.
Return type:: relax.Expr

tvm.script.relax.zeros(shape: tuple[int | Expr] | Expr, dtype: str | dtype, *, ty=None, span=None) → Expr#

Construct a tensor of all zeros, with the input shape and dtype.

Parameters:

shape (Union[Tuple[PrimExprLike], Expr]) – The shape of the created tensor.
dtype (str | DataType) – The data type of the created tensor.

Returns:

result – The result tensor.

Return type:

tvm.script.relax.zeros_like(x: Expr, dtype: str | dtype | None = None, *, ty=None, span=None) → Expr#

Construct a tensor with all zeros, with shape of the input tensor shape.

Parameters:

x (relax.Expr) – The input tensor, which provides the shape, and dtype when the dtype field is not specified.
dtype (Optional[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:

result – The result tensor.

Return type:

tvm.script.relax.function(function: LambdaType | None = None, **options: Any) → Any#

Parse a Python function into a function of the selected IR language variant.

Parameters:

function (Callable, optional) – The function to be parsed. May be omitted to use the decorator with keyword options, such as @T.function(private=True).
private (bool, optional) – Whether the function should be treated as private. A private function has no global symbol attribute; a public function has a global symbol matching its name. Defaults to False.
check_well_formed (bool, optional) – Whether to check that the constructed function is well formed. Defaults to True.
persistent (bool, optional) – For T.function, mark the resulting function as a persistent kernel. Defaults to False. See tvm.tirx.script.ir_builder.function().
pure (bool, optional) – For R.function, declare whether the function is pure, meaning that it has no observable side effects. Defaults to True. See tvm.relax.script.ir_builder.function_().
**options – Keyword options are forwarded to the selected language variant’s function_() hook, except check_well_formed, which controls parser validation. Options supported by only one language variant are not shared between T.function and R.function.

Returns:

result – The parsed function, or a decorator when function is omitted. Class members retain their Python functions until the enclosing module is constructed.

Return type:

Function or relax.Function or Callable

tvm.script.relax.macro(function: LambdaType | None = None, **options: Any) → Callable[[...], Any]#

Decorate a helper that constructs IR in its caller’s active frames.

Parameters:

function (Callable, optional) – The helper function. May be omitted to supply keyword options.
hygienic (bool, optional) – Whether the helper resolves symbols in its definition environment instead of its calling environment. Defaults to True. T.macro and R.macro capture values at definition time; T.inline refreshes captured closure cells when called.
**options – Keyword configuration for this helper decorator. hygienic controls name lookup as described above; other options are accepted but are not consumed or forwarded to builder hooks.

Returns:

result – The construction helper, or a decorator when function is omitted.

Return type:

Callable

Notes

T.inline follows Python lexical scoping with late binding of captured closure cells. Its return statements produce Python values, as do those of R.macro. T.macro emits returns in the active primitive function.

Examples

An inline helper can read values from its enclosing scope:

import tvm
from tvm.script import tirx as T
x_value = 128

@T.inline
def capture(A, B):
    B[()] = A[x_value]  # x_value resolved from enclosing scope

@T.function
def use(A: T.Tensor((1024,), "int32"), B: T.Tensor((), "int32")) -> None:
    capture(A, B)       # Produces B[()] = A[128]

tvm.script.tirx#

Operator attributes re-exported by this namespace are documented in tvm.tirx.

Public canonical TVMScript dialect namespace.

tvm.script.tirx.Lambda(parameter_types, function, *, ret_type=None)#

Build a shared staging lambda from explicit types and a Python callable.

Scalar constructors such as T.float32 may be used as parameter types. An optional return annotation checks the body type without inserting casts.

tvm.script.tirx.Let(expr: Expr, where: dict[Var, Expr]) → Expr#: Create a Let expression binding

tvm.script.tirx.Optional#: alias of OptionalAnnotation

tvm.script.tirx.Ptr(dtype, storage_scope='global', *, span=None)#

The pointer declaration function.

Parameters:

dtype (str, Type or callable) – The data type of the pointer.
storage_scope (str) – The storage scope of the pointer.
span (SpanEntry, Span or None, optional) – Source location attached to the constructed IR.

Returns:

res – The pointer.

Return type:

tvm.script.tirx.Tensor#: alias of _tensor_type

tvm.script.tirx.TensorLoad(buffer: Var, indices: list[Expr], span: Span | None = None) → TensorLoad#

Construct a validated buffer load.

Parameters:

buffer (Var) – The buffer to be loaded.
indices (List[Expr]) – The buffer indices to load values from.
span (Optional[Span]) – The location of this expression in the source code.

tvm.script.tirx.TensorMap() → Var#

Create a TIRx var that represents a CUDA tensor-map descriptor.

The host/runtime ABI passes a handle to descriptor storage. CUDA kernel codegen lowers this type to const __grid_constant__ CUtensorMap when it appears as a kernel parameter.

tvm.script.tirx.Tuple(*fields: Type) → Type#: Construct a tuple type for a TIRx function or binding annotation.

tvm.script.tirx.abi_field_get(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to tirx.abi_field_get.

tvm.script.tirx.abi_field_set(arr, index, field, value, *, ty=None, span=None)#

Set value in struct field in array

Parameters:

arr (StructType*) – The array of struct.
index (int) – The index of struct.
field (int) – The field of struct.
value (Expr) – The value to be set in field.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.abs(x, span=None)#

Get absolute value of the input element-wise.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

tvm.script.tirx.access_ptr(ptype, data, offset, extent, rw_mask, *, ty=None, span=None)#

Get head access address with memory access pattern info

Parameters:

ptype (PrimType or str) – The accessed element type. Offsets and extents are in units of this type.
data (DType*) – The data of pointer.
offset (int) – The offset of pointer.
extent (int) – The extent of pointer.
rw_mask (int) – The read write mask.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.acos(x, *, ty=None, span=None)#

Take acos of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.acosh(x, *, ty=None, span=None)#

Take acos of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.add_to_parent(stmt: Stmt) → None#: Add a statement to the parent frame.

tvm.script.tirx.address_of(obj: Var | TensorLoad, span: Span | None = None, *, ty=None) → Expr#

Returns the address of a buffer element or addressable variable.

Parameters:

obj (Union[Var, TensorLoad]) – The buffer, buffer load, or addressable variable.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.alloc_cast_frag(src, dtype)#

Allocate a register frag holding src value-cast to dtype.

Inherits src’s logical shape and its (lane, register) layout — only the element dtype changes — so Tx.cast(dst, src) is a per-thread element-wise cast with no cross-lane movement. .permute(...) the result to the axis order a downstream consumer (e.g. stmatrix via Tx.copy(dispatch="ldstmatrix")) expects.

Parameters:

src (Var) – Source register frag (e.g. from alloc_tcgen05_ldst_frag).
dtype (str) – Destination element dtype.

Returns:

Fresh local frag, src.shape shaped, src.layout, dtype-cast.

Return type:

tvm.script.tirx.alloc_local(shape: list[~tvm.ir.expr.Expr] | tuple[~tvm.ir.expr.Expr] | ~tvm.ir.expr.Expr | ~numbers.Integral, dtype: str = 'float32', data: ~tvm.ir.expr.Var | None = None, strides: list[~tvm.ir.expr.Expr] | None = None, elem_offset: ~tvm.ir.expr.Expr | None = None, byte_offset: ~tvm.ir.expr.Expr | None = None, *, scope: str = 'local', align: int = -1, offset_factor: int = 0, layout: str | ~tvm.tirx.layout.Layout | None = <MISSING>, allocated_addr: int | tuple[int, ...] | None = None, annotations: dict[str, ~typing.Any] | None = None) → Var#

Allocate a tensor and return its variable.

Emits a Bind statement with a tirx.alloc_tensor Call:

buf = T.alloc_tensor((128, 128))

Parameters:

shape (Union[List[Expr], Tuple[Expr], Expr, Integral]) – The shape of the buffer to allocate.
dtype (str) – The data type of the buffer elements.
scope (str) – The storage scope of the buffer (e.g., “global”, “shared”).
data (Optional[Var]) – Optional explicit data pointer.
strides (Optional[List[Expr]]) – Optional strides.
elem_offset (Optional[Expr]) – Optional element offset.
byte_offset (Optional[Expr]) – Optional byte offset.
align (int) – Alignment requirement in bytes.
offset_factor (int) – Offset factor.
layout (Optional[Union[str, Layout]]) – Optional layout.
allocated_addr (Optional[Union[int, Tuple[int, ...]]]) – Optional pre-allocated address metadata.
annotations (Optional[Dict[str, Any]]) – Optional annotations for the allocation.

Returns:

res – The allocated buffer.

Return type:

tvm.script.tirx.alloc_scalar(dtype: str = 'float32', scope: str = 'global', *, annotations: dict[str, Any] | None = None) → TensorLoad#: Allocate a zero-dimensional buffer (scalar), with optional allocation annotations.

tvm.script.tirx.alloc_shared(shape: list[~tvm.ir.expr.Expr] | tuple[~tvm.ir.expr.Expr] | ~tvm.ir.expr.Expr | ~numbers.Integral, dtype: str = 'float32', data: ~tvm.ir.expr.Var | None = None, strides: list[~tvm.ir.expr.Expr] | None = None, elem_offset: ~tvm.ir.expr.Expr | None = None, byte_offset: ~tvm.ir.expr.Expr | None = None, *, scope: str = 'shared', align: int = -1, offset_factor: int = 0, layout: str | ~tvm.tirx.layout.Layout | None = <MISSING>, allocated_addr: int | tuple[int, ...] | None = None, annotations: dict[str, ~typing.Any] | None = None) → Var#

Allocate a tensor and return its variable.

Emits a Bind statement with a tirx.alloc_tensor Call:

buf = T.alloc_tensor((128, 128))

Parameters:

shape (Union[List[Expr], Tuple[Expr], Expr, Integral]) – The shape of the buffer to allocate.
dtype (str) – The data type of the buffer elements.
scope (str) – The storage scope of the buffer (e.g., “global”, “shared”).
data (Optional[Var]) – Optional explicit data pointer.
strides (Optional[List[Expr]]) – Optional strides.
elem_offset (Optional[Expr]) – Optional element offset.
byte_offset (Optional[Expr]) – Optional byte offset.
align (int) – Alignment requirement in bytes.
offset_factor (int) – Offset factor.
layout (Optional[Union[str, Layout]]) – Optional layout.
allocated_addr (Optional[Union[int, Tuple[int, ...]]]) – Optional pre-allocated address metadata.
annotations (Optional[Dict[str, Any]]) – Optional annotations for the allocation.

Returns:

res – The allocated buffer.

Return type:

tvm.script.tirx.alloc_tcgen05_ldst_frag(instr_shape, tensor_shape, dtype)#

Allocate a register fragment for tcgen05.{ld,st} atoms.

Sizes the per-thread storage, allocates local scope memory, and returns a 2-D view of shape tensor_shape with a matching tcgen05_atom_layout. Pass the result to Tx.wg.copy_async (with a matching TMEM buffer) to trigger the corresponding dispatch path.

Parameters:

instr_shape (str) – "32x32b" (M=128 fragment, 128 row warpgroup tile, layout (128, K):(1@tid_in_wg, 1)); or "16x64b" / "16x128b" / "16x256b" (M=64 fragments, 64 row warpgroup tile with the per-shape per-lane register decomposition).
tensor_shape (tuple[int, int]) – Logical fragment shape (frag_rows, K) in element units. frag_rows is 128 for .32x32b and 64 for the .16x*b shapes. The fp32 Layout B readback image also uses ("32x32b", (64, N)).
dtype (str) – "float32", "float16", or "bfloat16".

Returns:

2-D view of shape tensor_shape whose layout matches tcgen05_atom_layout(instr_shape, tensor_shape, dtype).

Return type:

Examples

M=128 readback (existing dispatch):: frag = T.alloc_tcgen05_ldst_frag("32x32b", (128, 64), "float32") Tx.wg.copy_async(frag[:, :], tmem[:, 0:64])
M=64 readback (.16x64b dispatch):: frag = T.alloc_tcgen05_ldst_frag("16x64b", (64, 64), "float32") Tx.wg.copy_async(frag[:, :], tmem[0:64, 0:64])
Datapath B readback (cta_group=2, per-CTA M=64):: C = tmem_pool.alloc((64, 128), "float32", datapath="B") frag = T.alloc_tcgen05_ldst_frag("32x32b", (64, 128), "float32") Tx.wg.copy_async(frag[:, :], C[:, :])

tvm.script.tirx.alloc_tensor(shape: list[~tvm.ir.expr.Expr] | tuple[~tvm.ir.expr.Expr] | ~tvm.ir.expr.Expr | ~numbers.Integral, dtype: str = 'float32', data: ~tvm.ir.expr.Var | None = None, strides: list[~tvm.ir.expr.Expr] | None = None, elem_offset: ~tvm.ir.expr.Expr | None = None, byte_offset: ~tvm.ir.expr.Expr | None = None, scope: str = 'global', align: int = -1, offset_factor: int = 0, layout: str | ~tvm.tirx.layout.Layout | None = <MISSING>, allocated_addr: int | tuple[int, ...] | None = None, annotations: dict[str, ~typing.Any] | None = None) → Var#

Allocate a tensor and return its variable.

Emits a Bind statement with a tirx.alloc_tensor Call:

buf = T.alloc_tensor((128, 128))

Parameters:

shape (Union[List[Expr], Tuple[Expr], Expr, Integral]) – The shape of the buffer to allocate.
dtype (str) – The data type of the buffer elements.
scope (str) – The storage scope of the buffer (e.g., “global”, “shared”).
data (Optional[Var]) – Optional explicit data pointer.
strides (Optional[List[Expr]]) – Optional strides.
elem_offset (Optional[Expr]) – Optional element offset.
byte_offset (Optional[Expr]) – Optional byte offset.
align (int) – Alignment requirement in bytes.
offset_factor (int) – Offset factor.
layout (Optional[Union[str, Layout]]) – Optional layout.
allocated_addr (Optional[Union[int, Tuple[int, ...]]]) – Optional pre-allocated address metadata.
annotations (Optional[Dict[str, Any]]) – Optional annotations for the allocation.

Returns:

res – The allocated buffer.

Return type:

tvm.script.tirx.alloc_workspace(device_type, device_id, nbytes, dtype_code_hint, dtype_bits_hint, *, ty=None, span=None)#

Backend function to allocate temporal workspace

Parameters:

device_type (int) – The device type which the space will be allocated.
device_id (int) – The device id which the space will be allocated.
nbytes (int) – The size of the space requested.
dtype_code_hint (int) – The type code of the array elements. Only used in certain backends such as OpenGL.
dtype_bits_hint (int) – The type bits of the array elements. Only used in certain backends such as OpenGL.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.and_(*values: Any) → Any#: Implements tvm.script.ir_builder.parser_protocol.and_().

tvm.script.tirx.arg_(name: str, annotation: Any, *, span: SpanEntry | Span | None = None) → Var#: Implements tvm.script.ir_builder.parser_protocol.arg_().

tvm.script.tirx.asin(x, *, ty=None, span=None)#

Take asin of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.asinh(x, *, ty=None, span=None)#

Take asinh of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.assert_(condition: Any, message: str | tuple[str, Sequence[Any]] | Sequence[Any] = '', *, span: SpanEntry | Span | None = None) → None#: Implements tvm.script.ir_builder.parser_protocol.assert_().

tvm.script.tirx.assume(cond=None, *, ty=None, span=None)#

Provide a true statement that can be used for simplifications

Parameters:: cond (Expr) – The constraint condition.
Returns:: call – The call expression.
Return type:: Expr

tvm.script.tirx.assume_aligned(tensor, alignment_bytes, *, ty=None, span=None)#

Assume the tensor’s base address is aligned to alignment_bytes.

This compiler fact does not check or modify the address. The tensor must be a tensor variable and alignment a scalar integer constant, a power of two between 1 and 2**27 bytes (inclusive).

tvm.script.tirx.atan(x, *, ty=None, span=None)#

Take atan of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.atan2(x1, x2, *, ty=None, span=None)#

Take arctan2(x1, x2).

Parameters:

x1 (Expr) – Input argument.
x2 (Expr) – Input argument.

Returns:

y – The result.

Return type:

tvm.script.tirx.atanh(x, *, ty=None, span=None)#

Take atanh of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.bind(value: Expr, type_annotation: Type | None = None, *, var: Var | None = None) → Var#

Create a Bind (variable binding).

Emits a flat Bind statement to the current frame and returns the bound variable.

Parameters:

value (Expr) – The value to be bound.
type_annotation (Optional[Type] = None) – The type annotation of the binding. Usually it is used for fine-grained var typing, particularly, PointerType.
var (Optional[Var] = None) – The variable to bind. If not specified, a new variable will be created.

Returns:

var – The bound variable.

Return type:

tvm.script.tirx.bind_(value: ~typing.Any = <MISSING>, *, ty: ~typing.Any | None = None, name: str | None = None, span: ~tvm.script.ir_builder.base.SpanEntry | ~tvm.ir.base.Span | None = None, value_span: ~tvm.script.ir_builder.base.SpanEntry | ~tvm.ir.base.Span | None = None, name_span: ~tvm.script.ir_builder.base.SpanEntry | ~tvm.ir.base.Span | None = None, frame_value: bool = False) → Any#

Implements tvm.script.ir_builder.parser_protocol.bind_().

Returned Vars, including buffers, and metadata retain identity, names and spans. Other expressions create native Bind nodes; value_span belongs to the RHS. Explicit typed bindings and frame targets retain their separate contracts.

tvm.script.tirx.bitwise_and(x, y, span=None)#

Take bitwise and of two values

Parameters:

x (Expr) – Left operand
y (Expr) – Right operand
span (Optional[Span]) – The location of this operator in the source code.

Returns:

res – The result.

Return type:

tvm.script.tirx.bitwise_not(x, span=None)#

Take bitwise not of input value

Parameters:

x (Expr) – Input operand
span (Optional[Span]) – The location of this operator in the source code.

Returns:

res – The result.

Return type:

tvm.script.tirx.bitwise_or(x, y, span=None)#

Take bitwise or of two values

Parameters:

x (Expr) – Left operand
y (Expr) – Right operand
span (Optional[Span]) – The location of this operator in the source code.

Returns:

res – The result.

Return type:

tvm.script.tirx.bitwise_xor(x, y, span=None)#

Take bitwise xor of two values

Parameters:

x (Expr) – Left operand
y (Expr) – Right operand
span (Optional[Span]) – The location of this operator in the source code.

Returns:

res – The result.

Return type:

tvm.script.tirx.bool(expr: Expr | None = None) → Expr#

Construct a new tirx.Var with type boolean or cast expression to type boolean.

Parameters:: expr (Expr) – The expression to be cast.
Returns:: res – The new tirx.Var with type boolean or casted expression with type boolean.
Return type:: Expr

tvm.script.tirx.boolean(expr: Expr | None = None) → Expr#

Construct a new tirx.Var with type boolean or cast expression to type boolean.

Parameters:: expr (Expr) – The expression to be cast.
Returns:: res – The new tirx.Var with type boolean or casted expression with type boolean.
Return type:: Expr

tvm.script.tirx.break_(*, span: SpanEntry | Span | None = None) → AlreadyEmitted[Stmt]#

Implements tvm.script.ir_builder.parser_protocol.break_().

Legality is checked on the completed function, across loop and function boundaries.

tvm.script.tirx.broadcast#: alias of Broadcast

tvm.script.tirx.call_cpacked(*args, span=None, ty=None)#

Build expression by call an external packed function.

Same as call_packed, except that the first argument is the function name (as in call_extern), and the last argument is the resource handle.

Parameters:

args (list of Expr or Var.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

See also

te.extern: Create tensor with extern function call.

tvm.script.tirx.call_cpacked_lowered(*args, span=None, ty=None)#

Lowered version of call c-packed. Same as call_packed, except that the first argument is the function name (as in call_extern), and the last argument is the resource handle.

Parameters:

args (list of Expr or Var.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

See also

te.extern: Create tensor with extern function call.

tvm.script.tirx.call_extern(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to tirx.call_extern.

tvm.script.tirx.call_ffi_kernel(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to tirx.call_ffi_kernel.

tvm.script.tirx.call_global_var_(function: GlobalVar, args: Sequence[Any]) → Expr#: Implements tvm.script.ir_builder.parser_protocol.call_global_var_().

tvm.script.tirx.call_intrin(func_name, *args, ty, attrs=None, span=None)#: Call an intrinsic with an explicit result type.

tvm.script.tirx.call_kernel(kernel, launch_args: list[int | Expr | list[int | Expr]], *args: list[Any], **kwargs: dict[str, Any])#

Call an external kernel.

Parameters:

kernel (Any) – The external kernel to call.
launch_args (List[Union[int, tirx.Expr, List[Union[int, tirx.Expr]]]]) – The launch arguments. A list of integers for grid size, block size, and shared memory size. The actual requirements depend on the kernel.
args (List[tirx.Expr]) – The arguments to pass to the kernel.
kwargs (Dict[str, Any]) – Additional keyword arguments to pass to the kernel or compilation.

tvm.script.tirx.call_llvm_intrin(dtype, name, *args, span=None)#

Build expression by calling a llvm intrinsic function

Parameters:

dtype (str) – The data type of the result.
name (str) – The name of the llvm intrinsic function.
args (list) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.call_llvm_pure_intrin(dtype, name, *args, span=None)#

Build expression by calling a pure llvm intrinsic function

Parameters:

dtype (str) – The data type of the result.
name (str) – The name of the llvm intrinsic function.
args (list) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.call_packed(*args, span=None, ty=None)#

Build expression by call an external packed function.

The argument to a packed function can be an Expr or a tensor variable. The argument is the corresponding POD type when Expr is presented.

When the argument is a Var carrying TensorType, the corresponding PackedFunc will receive an TVMArrayHandle whose content is valid during the callback period. If the PackedFunc is a python callback, then the corresponding argument is Tensor.

Parameters:

args (list of Expr or Var.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

See also

te.extern: Create tensor with extern function call.

tvm.script.tirx.call_packed_lowered(*args, span=None, ty=None)#

Lowered version of call packed. The argument to a packed function can be an Expr or a tensor variable. The argument is the corresponding POD type when Expr is presented. When the argument is a Var carrying TensorType, the corresponding PackedFunc will receive an TVMArrayHandle whose content is valid during the callback period. If the PackedFunc is a python callback, then the corresponding argument is Tensor.

Parameters:

args (list of Expr or Var.) – Positional arguments.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

call – The call expression.

Return type:

See also

te.extern: Create tensor with extern function call.

tvm.script.tirx.call_pure_extern(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to tirx.call_pure_extern.

tvm.script.tirx.cast(value, dtype, span=None)#: Cast an expression to the requested data type.

tvm.script.tirx.ceil(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to prim.ceil.

tvm.script.tirx.ceildiv(lhs, rhs, span=None)#

Generic ceildiv operator.

Parameters:

lhs (object) – The left operand.
rhs (object) – The right operand.
span (Optional[Span]) – The location of this operator in the source.

Returns:

op – The result Expr of ceildiv operaton.

Return type:

tvm.Expr

tvm.script.tirx.check_well_formed_(function: Function) → None#: Validate a completed TIRx function.

tvm.script.tirx.cluster_id(extents: list[Expr | int] | None = None, dtype: str = 'int32') → Var | tuple[Var, ...]#

Define a kernel→cluster scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

tvm.script.tirx.clz(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to prim.clz.

tvm.script.tirx.compute_scope(name: str) → RegionFrame#: Outline the body as a named CPU compute helper.

tvm.script.tirx.constexpr(value: object) → NoReturn#

Mark host control syntax or a JIT specialization annotation.

Parameters:: value (object) – Source expression to evaluate with ordinary Python semantics in a supported control-flow position. The marker itself can also appear as an annotation identifying a JIT specialization parameter.
Raises:: TypeError – If invoked directly instead of being recognized in parsed source.

Notes

# Source
if I.constexpr(enabled):
    T.evaluate(1)
# Builder
if enabled:
    X.emit_(X.evaluate(1))

tvm.script.tirx.continue_(*, span: SpanEntry | Span | None = None) → AlreadyEmitted[Stmt]#

Implements tvm.script.ir_builder.parser_protocol.continue_().

Legality is checked on the completed function, across loop and function boundaries.

tvm.script.tirx.copysign(x1, x2, *, ty=None, span=None)#

Change the sign of x1 to that of x2, element-wise.

Parameters:

x1 (Expr) – Input argument.
x2 (Expr) – Input argument.

Returns:

y – The result.

Return type:

tvm.script.tirx.cos(x, *, ty=None, span=None)#

Take cos of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.cosh(x, *, ty=None, span=None)#

Take cosh of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.cpu_parallel_barrier(*, ty=None, span=None)#

Synchronize all workers in the current CPU parallel launch.

Every worker must reach this operation at the same program point. Place it inside a parallel_launch region, outside parallel loops, whose iteration counts can differ between workers. Writes before the barrier are visible to all workers after it.

To replace pragma_parallel_barrier_when_finish, place this operation after the former attribute body.

tvm.script.tirx.cta_id(extents: list[Expr | int] | None = None, preferred=None, dtype: str = 'int32') → Var | tuple[Var, ...]#

Define a kernel→cta scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

tvm.script.tirx.cta_id_in_cluster(extents: list[Expr | int] | None = None, preferred=None, dtype: str = 'int32') → Var | tuple[Var, ...]#

Define a cluster→cta scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

tvm.script.tirx.cta_id_in_pair(dtype: str = 'int32') → Var#: Return the native CTA index within its two-CTA pair.

tvm.script.tirx.decl_mutable_cell_(value: ~typing.Any = <MISSING>, *, ty: ~typing.Any | None = None, name: str | None = None, span: ~tvm.script.ir_builder.base.SpanEntry | ~tvm.ir.base.Span | None = None, name_span: ~tvm.script.ir_builder.base.SpanEntry | ~tvm.ir.base.Span | None = None) → Any#

Implements tvm.script.ir_builder.parser_protocol.decl_mutable_cell_().

Primitive annotations allocate scalar local storage; vector annotations allocate their declared shape. Var declaration producers retain their own effects.

tvm.script.tirx.decl_scalar(dtype, data, scope, elem_offset=None, byte_offset=None) → TensorLoad#: Declare a zero-dimensional buffer (scalar) from a pointer.

tvm.script.tirx.decl_tensor(shape, dtype='float32', data=None, strides=None, elem_offset=None, byte_offset=None, scope='global', align=0, offset_factor=0, layout=<MISSING>, allocated_addr=None) → Var#

Declare a tensor backed by a pointer, or allocate its storage.

With data, bind a tirx.decl_tensor Call to the tensor variable. Without data, bind a tirx.alloc_tensor Call instead. The tmem scope uses allocated_addr to declare externally allocated tensor memory.

Parameters:

shape (Union[List[Expr], Tuple[Expr], Expr, Integral]) – The type of the buffer prior to flattening.
dtype (str) – The data type in the content of the buffer.
data (Var) – The pointer to the head of the data.
strides (List[Expr]) – The strides of each dimension.
elem_offset (Expr) – The offset in terms of number of dtype elements (including lanes).
byte_offset (Expr) – The offset in terms of number of bytes.
scope (str) – The optional storage scope of buffer data pointer.
align (int) – The alignment requirement of data pointer in bytes.
offset_factor (int) – The factor of elem_offset field.
layout (Layout) – The layout of the buffer.

Returns:

res – The declared buffer.

Return type:

tvm.script.tirx.device_context(device_type: Expr, device_id: Expr) → RegionFrame#

Supply lexical device context for allocation and packed-call lowering.

This region does not change the active runtime device.

tvm.script.tirx.device_entry() → RegionFrame#

Mark a device-entry region containing scope definitions.

Use a flat T.device_entry() to scope the remaining statements in the enclosing body, or with T.device_entry(): for an explicit boundary. Statements before the region remain host code.

tvm.script.tirx.dtype(value: str | dtype) → DataTypeImm#

Construct a DataType-valued constant for an ordinary operation argument.

Unlike scalar constructors such as T.float32, this constructs a dtype value, rather than a scalar value or variable of that dtype.

tvm.script.tirx.dynamic(name: str, dtype: str = 'int64', *, span: SpanEntry | Span | None = None) → Var#

Create a fresh primitive symbolic variable, independently of builder scope.

Parameters:

name (str) – The symbol’s display name. Repeated names do not share identity.
dtype (str) – Primitive dtype, defaulting to int64.
span (SpanEntry, Span or None) – Source location of the symbol.

Returns:

A fresh symbol. Reuse this object to share dimensions across annotations and function bodies, including outside an I.ir_module definition.

Return type:

tvm.script.tirx.else_(*, span: SpanEntry | Span | None = None) → ElseFrame#: Implements tvm.script.ir_builder.parser_protocol.else_().

tvm.script.tirx.emit(value: Any, *, span: SpanEntry | Span | None = None) → None#

Implements tvm.script.ir_builder.parser_protocol.emit_().

Native statements emit once; receipts are already emitted. Vars, layouts and meta_class instances are inert. Sequences are consumed elementwise; concise frames close with their owning parent. Other values use native expression conversion, retaining its errors for unsupported host values.

tvm.script.tirx.emit_(value: Any, *, span: SpanEntry | Span | None = None) → None#

Implements tvm.script.ir_builder.parser_protocol.emit_().

tvm.script.tirx.eq_(lhs: Any, rhs: Any, *, span: SpanEntry | Span | None = None) → Expr#: Implements tvm.script.ir_builder.parser_protocol.eq_().

tvm.script.tirx.erf(x, *, ty=None, span=None)#

Take gauss error function of the input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.evaluate(value: Expr) → AlreadyEmitted[Stmt]#

Emit an evaluation and return a reference to its stored statement.

Parameters:: value (Expr) – The input expression to evaluate.
Returns:: result – A receipt containing the emitted statement, so expression-statement handling does not emit it again.
Return type:: AlreadyEmitted[Stmt]

tvm.script.tirx.exp(x, *, ty=None, span=None)#

Take exponential of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.exp10(x, *, ty=None, span=None)#

Calculate 10**x

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.exp2(x, *, ty=None, span=None)#

Calculate 2**x

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.fabs(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to prim.fabs.

tvm.script.tirx.floor(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to prim.floor.

tvm.script.tirx.floordiv(a, b, span=None)#

Compute the floordiv of two expressions.

Parameters:

a (Expr) – The left hand operand
b (Expr) – The right hand operand
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

tvm.script.tirx.floormod(a, b, span=None)#

Compute the floormod of two expressions.

Parameters:

a (Expr) – The left hand operand
b (Expr) – The right hand operand
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

tvm.script.tirx.fma(x, y, z, *, ty=None, span=None)#

Take fused multiply-add of input x, y, z.

Parameters:

x (Expr) – First input argument.
y (Expr) – Second input argument.
z (Expr) – Third input argument.

Returns:

out – The result of x * y + z.

Return type:

tvm.script.tirx.fmod(x, y, *, ty=None, span=None)#

Return the remainder of x divided by y with the same sign as x.

Parameters:

x (Expr) – Input argument.
y (Expr) – Input argument.

Returns:

z – The result.

Return type:

tvm.script.tirx.for_(iterable: Any, *, names: str | Sequence[str] | None = None, span: SpanEntry | Span | None = None) → ForFrame#

Implements tvm.script.ir_builder.parser_protocol.for_().

A single native loop returns its scalar Var; multiple loops return their sequence. frame.vars remains the stable sequence for source unpacking.

tvm.script.tirx.frame_scope(frames: list[StmtFrame]) → _FrameScope#

Enter multiple IRBuilder frames without deep nesting.

This function provides a way to enter multiple frames in a single with statement, which is particularly useful when migrating from cases where allocations don’t require nested scopes.

Parameters:: frames (List[frame.StmtFrame]) – The list of frames to enter. Each frame’s __enter__ return value will be collected and returned as a tuple.
Returns:: A context manager that enters all frames and returns their values.
Return type:: _FrameScope

tvm.script.tirx.free_workspace(device_type, device_id, ptr, *, ty=None, span=None)#

Backend function to free temporal workspace.

Parameters:

device_type (int) – The device type which the space will be allocated.
device_id (int) – The device id which the space will be allocated.
ptr (Var) – The result allocated space pointer.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.func_attr(attrs: dict[str, Any]) → None#

The Function annotation statement.

Parameters:: attrs (Dict[str, Any]) – The annotations of the Function.

tvm.script.tirx.func_name_(name: str) → None#: Implements tvm.script.ir_builder.parser_protocol.func_name_().

tvm.script.tirx.func_ret(annotation: Any, *, span: SpanEntry | Span | None = None) → None#: Implements tvm.script.ir_builder.parser_protocol.func_ret_type_().

tvm.script.tirx.func_ret_type_(annotation: Any, *, span: SpanEntry | Span | None = None) → None#: Implements tvm.script.ir_builder.parser_protocol.func_ret_type_().

tvm.script.tirx.function_(*, private: bool = False, persistent: bool = False, decl: bool = False, span: SpanEntry | Span | None = None) → FunctionFrame#

Implements tvm.script.ir_builder.parser_protocol.function_().

Private/persistent options pass to the native TIRx function frame. The same frame supports declaration and body entry.

tvm.script.tirx.ge_(lhs: Any, rhs: Any, *, span: SpanEntry | Span | None = None) → Expr#: Implements tvm.script.ir_builder.parser_protocol.ge_().

tvm.script.tirx.get_active_lane_mask(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to tirx.get_active_lane_mask.

tvm.script.tirx.gpu_active_thread_selector(var, pred, span=None, *, ty=None)#

Analysis-only active-thread selector.

gpu_active_thread_selector(var, pred) denotes the unique value of var in the current active domain for which pred is true. It is intended for compiler metadata and should not survive to executable codegen.

tvm.script.tirx.gpu_dp4a(vec1, vec2, acc=0, **kwargs)#

Dot product of two int8x4 vectors and add an optional accumulator

Parameters:

vec1 (int8x4) – The input vector.
vec2 (int8x4) – The input vector.
acc (int32) – The accumulator.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.gpu_fill_fragment(fragment, m, n, k, index, value, *, ty=None, span=None)#

TVM intrinsic for tensor core fill_fragment operators

Parameters:

fragment (Var) – The wmma fragment
m (UIntImm) – The shape of wmma fragment.
n (UIntImm) – The shape of wmma fragment.
k (UIntImm) – The shape of wmma fragment.
index (Expr) – The fragment index.
value (Expr) – The value to be filled in fragment.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.gpu_load_matrix_sync(fragment, m, n, k, index, buffer_ptr, stride, layout, *, ty=None, span=None)#

TVM intrinsic for tensor core load operators

Parameters:

fragment (Var) – The wmma fragment.
m (UIntImm) – The shape of wmma fragment.
n (UIntImm) – The shape of wmma fragment.
k (UIntImm) – The shape of wmma fragment.
index (Expr) – The fragment index.
buffer_ptr (Expr) – The fragment buffer pointer.
stride (Expr) – The fragment stride.
layout (Literal["row_major", "column_major"]) – The fragment layout.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.gpu_mma_sync(fragment_d, index_d, fragment_a, index_a, fragment_b, index_b, fragment_c, index_c, *, ty=None, span=None)#

TVM intrinsic for tensor core mma_sync operators

Parameters:

fragment_d (Var) – The wmma fragment_d.
index_d (Expr) – The fragment_d index.
fragment_a (Var) – The wmma fragment_a.
index_a (Expr) – The fragment_a index.
fragment_b (Var) – The wmma fragment_b.
index_b (Expr) – The fragment_b index.
fragment_c (Var) – The wmma fragment_c.
index_c (Expr) – The fragment_c index.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.gpu_storage_sync(storage_scope, *, ty=None, span=None)#

Synchronize accesses in the specified storage scope.

Parameters:

storage_scope (str) – The storage scope to synchronize.
ty (tvm.ir.Type or str, optional) – The result type. Inferred as void when omitted.
span (tvm.ir.Span, optional) – The source location.

tvm.script.tirx.gpu_store_matrix_sync(fragment, m, n, k, index, buffer_ptr, stride, layout, *, ty=None, span=None)#

TVM intrinsic for tensor core store operators

Parameters:

fragment (Var) – The wmma fragment.
m (UIntImm) – The shape of wmma fragment.
n (UIntImm) – The shape of wmma fragment.
k (UIntImm) – The shape of wmma fragment.
index (Expr) – The fragment index.
buffer_ptr (Expr) – The fragment buffer pointer.
stride (Expr) – The fragment stride.
layout (Literal["row_major", "column_major"]) – The fragment layout.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.gpu_thread_allreduce(combine, identity, values, predicate, destinations, thread_axes, *, ty=None, span=None)#

Perform an all-reduce inside a thread block.

Parameters:

combine (tvm.ir.LambdaExpr) – Typed combining lambda with parameters ordered as all left-hand values followed by all right-hand values. Its body returns a scalar for one result or a Tuple of results.
identity (Expr or Sequence[Expr]) – Identity value for each reduction result.
values (Expr or Sequence[Expr]) – Values contributed by the current thread.
predicate (PrimExpr) – Boolean participation predicate. Inactive threads contribute identities.
destinations (Expr or Sequence[Expr]) – Tensor loads at index zero of one-element result temporaries, optionally cast for boolean storage. Each temporary must be accessed only at index zero.
thread_axes (Expr or Sequence[Expr]) – Thread variables participating in the reduction.

Returns:

call – The void call expression with six explicit operands.

Return type:

tvm.script.tirx.gpu_thread_filter(var, pred, *, span=None, ty=None)#

Thread-set filter escape hatch.

Use this wrapper only when the predicate is not in the canonical thread-filter grammar (see src/tirx/analysis/filter_canonical.h). Canonical predicates – pure conjunctions of scopeid_var <op> const comparisons plus bare T.cuda.elect_sync() calls – are recognized by the lowering pass directly from if cond:, so the wrapper is redundant for them.

When wrapped: var (a ScopeIdDef-declared scope identifier) tells the compiler which active-set axis to collapse to a singleton when the opaque predicate evaluates true; pred is preserved verbatim and evaluated at runtime.

tvm.script.tirx.gpu_thread_invariant(cond, *, ty=None, span=None)#

Mark condition as thread invariant.

Parameters:: cond (Expr) – The condition.
Returns:: call – The call expression.
Return type:: Expr

tvm.script.tirx.gpu_thread_return(*, ty=None, span=None)#: Return from the current GPU thread without a function value.

tvm.script.tirx.gpu_warp_activemask(*, ty=None, span=None)#

Return a 32-bit mask indicates currently active threads in a calling warp.

Returns:: call – The call expression.
Return type:: Expr

tvm.script.tirx.gpu_warp_shuffle(mask, value, warp_id, width, warp_size, *, ty=None, span=None)#

Exchange value between threads inside a warp.

Parameters:

mask (Expr) – The warp mask indicates active threads inside warp.
value (Expr) – The value to exchange.
warp_id (Expr) – The source lane index to fetch value.
width (Expr) – The width of sub-sections to perform warp shuffle.
warp_size (Expr) – The warp size.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.gpu_warp_shuffle_down(mask, value, offset, width, warp_size, *, ty=None, span=None)#

Copy value from a lane with higher (by offset) index relative to caller.

Parameters:

mask (Expr) – The warp mask indicates active threads inside warp.
value (Expr) – The value to exchange.
offset (Expr) – The difference between source lane index and destination lane index: offset = src_lane_idx - dst_lane_idx
width (Expr) – The width of sub-sections to perform warp shuffle.
warp_size (Expr) – The warp size.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.gpu_warp_shuffle_up(mask, value, offset, width, warp_size, *, ty=None, span=None)#

Copy value from a lane with lower (by offset) index relative to caller.

Parameters:

mask (Expr) – The warp mask indicates active threads inside warp.
value (Expr) – The value to exchange.
offset (Expr) – The difference between source lane index and destination lane index: offset = dst_lane_idx - src_lane_idx
width (Expr) – The width of sub-sections to perform warp shuffle.
warp_size (Expr) – The warp size.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.gpu_warp_shuffle_xor(mask, value, lane_mask, width, warp_size, *, ty=None, span=None)#

Copy value from a lane with index computed by src_lane_idx ^ lane_mask.

Parameters:

mask (Expr) – The warp mask indicates active threads inside warp.
value (Expr) – The value to exchange.
lane_mask (Expr) – The mask to compute source lane index:
width (Expr) – The width of sub-sections to perform warp shuffle.
warp_size (Expr) – The warp size.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.grid(*extents: tuple[Expr | tuple[Expr, Expr]], dtype: str | None = None) → ForFrame#

The grid For statement.

Parameters:

extents (Tuple[Union[Expr, Tuple[Expr, Expr]]]) – If a single Expr is provided, it is used as the extent of the iteration. If a tuple of two Expr is provided, the first is the start of the iteration, and the second is the extent of the iteration.
dtype (str, optional) – The dtype of every loop variable, either "int32" or "uint32". When omitted each loop variable takes the dtype of its own extent.

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

tvm.script.tirx.gt_(lhs: Any, rhs: Any, *, span: SpanEntry | Span | None = None) → Expr#: Implements tvm.script.ir_builder.parser_protocol.gt_().

tvm.script.tirx.handle(dtype: str | None = None, storage_scope: str = 'global') → Var#

Create a TIR var that represents a pointer.

Parameters:

dtype (str | None) – The data type of the pointer. If omitted, construct an opaque handle.
storage_scope (str) – The storage scope of the pointer.

Returns:

res – The new tirx.Var with type handle or casted expression with type handle.

Return type:

tvm.script.tirx.handle_add_byte_offset(handle, offset, *, ty=None, span=None)#

Add offset to handle

Parameters:

handle (Expr) – The handle.
offset (int) – The offset.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.hypot(x1, x2, *, ty=None, span=None)#

Equivalent to sqrt(x1**2 + x2**2), element-wise.

Parameters:

x1 (Expr) – Input argument.
x2 (Expr) – Input argument.

Returns:

y – The result.

Return type:

tvm.script.tirx.if_(condition: Any, *, span: SpanEntry | Span | None = None) → IfFrame#: Implements tvm.script.ir_builder.parser_protocol.if_().

tvm.script.tirx.if_then_else(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to prim.if_then_else.

tvm.script.tirx.if_then_else_(condition: Any, true_value: Any, false_value: Any) → Any#: Implements tvm.script.ir_builder.parser_protocol.if_then_else_().

tvm.script.tirx.ignore_loop_partition(predicate, *, ty=None, span=None) → Expr#

Annotate a predicate not be considered as target condition of loop partition.

Parameters:: predicate (Expr) – The annotated predicate expression.

tvm.script.tirx.index_map(mapping: Callable, *, inverse_index_map: Callable | None = None, index_dtype: str = 'int64') → IndexMap#: Create a TIR Index mapping

tvm.script.tirx.infinity(dtype: str, span: Span | None = None) → Any#

infinity value of dtype

Parameters:

dtype (str) – The data type.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

value – The infinity value of dtype.

Return type:

tvm.Expr

tvm.script.tirx.isfinite(x, span=None)#

Check if input value is finite.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

tvm.script.tirx.isinf(x, span=None)#

Check if input value is infinite.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

tvm.script.tirx.isnan(x, span=None, *, ty=None)#

Check if input value is Nan.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

tvm.script.tirx.isnullptr(x, span=None, *, ty=None)#

Check if input value is nullptr.

Parameters:

x (Expr) – Input argument.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

y – The result.

Return type:

tvm.script.tirx.kernel_replace_point(*, ty=None, span=None)#: Mark where a transform should replace generated kernel initialization.

tvm.script.tirx.lane_id(extents: list[Expr | int] | None = None, dtype: str = 'int32') → Var | tuple[Var, ...]#

Define a warp→thread scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

tvm.script.tirx.launch_thread(thread_tag: str, extent: Expr) → RegionFrame#

Launch a hardware or virtual thread with a fresh lexical variable.

The extent determines the variable’s scalar integer type. Tags starting with vthread denote virtual threads. This frame works with any TIR builder parent, including standalone statement construction.

Examples

with T.launch_thread("threadIdx.x", 32) as tx:
    T.evaluate(tx)

tvm.script.tirx.ldexp(x1, x2, *, ty=None, span=None)#

Returns x1 * (2 ** x2).

Parameters:

x1 (Expr) – Input argument.
x2 (Expr) – Input argument.

Returns:

y – The result.

Return type:

tvm.script.tirx.le_(lhs: Any, rhs: Any, *, span: SpanEntry | Span | None = None) → Expr#: Implements tvm.script.ir_builder.parser_protocol.le_().

tvm.script.tirx.likely(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to prim.likely.

tvm.script.tirx.llvm_lookup_intrinsic_id(name)#

Lookup LLVM intrinsic id by name.

Parameters:: name (str) – The name of the intrinsic.
Returns:: intrin_id – The intrinsic id.
Return type:: int

tvm.script.tirx.local_scalar(dtype: str = 'float32', *, annotations: dict[str, Any] | None = None) → TensorLoad#: Allocate a zero-dimensional buffer in local memory.

tvm.script.tirx.log(x, *, ty=None, span=None)#

Take log of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.log10(x, *, ty=None, span=None)#

Take log10 of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.log1p(x, *, ty=None, span=None)#

Take log(x + 1) with respect to input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.log2(x, *, ty=None, span=None)#

Take log2 of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.logical_and(*values)#

Construct scalar or vector conjunction from eager operands.

Parameters:: values (Expr or Python value) – One or more operands. Object-convertible values are normalized first. Host pairs follow Python logical operations; IR pairs use scalar logical or vector bitwise operations according to their types.
Returns:: result – The conjunction reduced from left to right.
Return type:: Expr or Python value

Notes

All arguments are evaluated before this call; it does not provide Python short-circuit evaluation of the argument expressions.

tvm.script.tirx.logical_not(value)#

Negate a host or IR value.

Parameters:: value (Expr or Python value) – Operand to negate. Object-convertible values are normalized first; IR expressions use primitive Not and host values use Python not.
Returns:: result – The logical negation without testing an IR expression as a Python bool.
Return type:: Expr or bool

tvm.script.tirx.logical_or(*values)#

Construct scalar or vector disjunction from eager operands.

Parameters:: values (Expr or Python value) – One or more operands. Object-convertible values are normalized first. Host pairs follow Python logical operations; IR pairs use scalar logical or vector bitwise operations according to their types.
Returns:: result – The disjunction reduced from left to right.
Return type:: Expr or Python value

Notes

All arguments are evaluated before this call; it does not provide Python short-circuit evaluation of the argument expressions.

tvm.script.tirx.lt_(lhs: Any, rhs: Any, *, span: SpanEntry | Span | None = None) → Expr#: Implements tvm.script.ir_builder.parser_protocol.lt_().

tvm.script.tirx.masked_load(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to tirx.masked_load.

tvm.script.tirx.masked_store(buffer, value, *indices_and_mask, ty=None, span=None)#

Store vector lanes selected by a predicate mask.

Parameters:

buffer (Var) – The buffer to update.
value (Expr) – The vector value to store.
indices_and_mask (Expr) – The buffer indices followed by a boolean lane mask. The mask must match the lane count and scalability of value.

Returns:

call – A void-typed tirx.masked_store call.

Return type:

tvm.script.tirx.max(a: Expr, b: Expr) → Expr#

Compute the maximum value of two expressions.

Parameters:

a (Expr) – The left hand operand
b (Expr) – The right hand operand

Returns:

res – The result expression.

Return type:

tvm.script.tirx.max_value(dtype: str, span: Span | None = None) → Any#

maximum value of dtype

Parameters:

dtype (str) – The data type.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

value – The maximum value of dtype.

Return type:

tvm.ir.Expr

tvm.script.tirx.meta_class(cls)#

Decorator for utility classes used inside @T.function.

Instances of decorated classes are treated as parser meta values.

tvm.script.tirx.min(a: Expr, b: Expr) → Expr#

Compute the minimum value of two expressions.

Parameters:

a (Expr) – The left hand operand
b (Expr) – The right hand operand

Returns:

res – The result expression.

Return type:

tvm.script.tirx.min_value(dtype, span=None)#

minimum value of dtype

Parameters:

dtype (str) – The data type.
span (Optional[Span]) – The location of this operator in the source code.

Returns:

value – The minimum value of dtype.

Return type:

tvm.ir.Expr

tvm.script.tirx.ne_(lhs: Any, rhs: Any, *, span: SpanEntry | Span | None = None) → Expr#: Implements tvm.script.ir_builder.parser_protocol.ne_().

tvm.script.tirx.nearbyint(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to prim.nearbyint.

tvm.script.tirx.nextafter(x1, x2, *, ty=None, span=None)#

Return the next floating-point value after x1 towards x2.

Parameters:

x1 (Expr) – Input argument.
x2 (Expr) – Input argument.

Returns:

y – The result.

Return type:

tvm.script.tirx.not_(value: Any) → Any#: Implements tvm.script.ir_builder.parser_protocol.not_().

tvm.script.tirx.or_(*values: Any) → Any#: Implements tvm.script.ir_builder.parser_protocol.or_().

tvm.script.tirx.parallel(start: Expr, stop: Expr | None = None, *, annotations: dict[str, Any] | None = None, step: Expr | None = None, dtype: str | None = None) → ForFrame#

The parallel For statement.

Parameters:

start (Expr) – The minimum value of iteration.
stop (Expr) – The maximum value of iteration.
annotations (Dict[str, Any]) – The optional annotations of the For statement. On CPU, {"parallel_stride_pattern": True} assigns iterations cyclically: worker t executes t, t + P, t + 2 * P, ... for P workers. Absent or false selects contiguous chunks. Each loop selects its own policy. When migrating pragma_parallel_stride_pattern, annotate every intended loop, including later loops in the same launch outside the former attribute body.
step (Expr) – The optional step value of iteration.
dtype (str, optional) – The dtype of the loop variable, either "int32" or "uint32". When omitted it is inferred from the bounds.

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

tvm.script.tirx.parallel_launch() → RegionFrame#: Launch a CPU worker team around parallel loops and team barriers.

tvm.script.tirx.popcount(x, *, ty=None, span=None)#

Count the number of set bits in input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.pow(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to prim.pow.

tvm.script.tirx.ptr(dtype: str, storage_scope: str = 'global') → Var#

The pointer declaration function.

Parameters:

dtype (str) – The data type of the pointer.
storage_scope (str) – The storage scope of the pointer.

Returns:

res – The pointer.

Return type:

tvm.script.tirx.ptr_byte_offset(data, byte_offset, dtype, *, ty=None, span=None)#

Cast data + byte_offset to dtype*.

byte_offset is always in bytes. Use this when the source CUDA shape needs an explicitly typed local pointer derived from a byte-addressed base.

tvm.script.tirx.ramp#: alias of Ramp

tvm.script.tirx.range_(*args: Any, annotations: dict[str, Any] | None = None) → ForFrame#: Implements tvm.script.ir_builder.parser_protocol.range_().

tvm.script.tirx.region(op: Op | str, args: Sequence[Expr], body_params: Sequence[Var] | None = None, attrs: DictAttrs | dict[str, Any] | None = None) → RegionFrame#

Construct a result-free region with operation-defined body parameters.

When body_params is omitted, the operation’s FRegionGetBodyParams hook creates fresh typed variables. Every region operation must register the hook, returning an empty array for no body parameters. Missing hooks reject construction even with explicit parameters. Explicit parameters must match the hook’s count and types and retain their identities.

Operands and attributes belong to the enclosing scope. Entering the frame returns one parameter directly, or a sequence for zero or multiple parameters. Result variables are supported by direct tvm.ir.RegionStmt construction; outward-result script syntax is not supported.

tvm.script.tirx.register_script_namespace(name: str, namespace: object) → object#

Parameters:

name (str) – Namespace name on the TIRx builder facade.
namespace (object) – Construction namespace object.

tvm.script.tirx.reinterpret(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to tirx.reinterpret.

tvm.script.tirx.resolve_global_info_(content: Any) → Any#

Implements tvm.script.ir_builder.parser_protocol.resolve_global_info_().

TIRx does not define global-info selectors.

tvm.script.tirx.return_(value: Any | None = None, *, span: SpanEntry | Span | None = None) → AlreadyEmitted[Stmt]#: Implements tvm.script.ir_builder.parser_protocol.return_().

tvm.script.tirx.round(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to prim.round.

tvm.script.tirx.rsqrt(x, *, ty=None, span=None)#

Take reciprocal of square root of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.scope_id(extents: list[Expr | int] | None, parent: str, cur: str, dtype: str = 'int32') → Var | tuple[Var, ...]#

Declare scope IDs between execution levels and return their native variables.

One dimension returns a variable; multiple dimensions return a tuple in declaration order. None defers extent inference to LowerTIRx.

tvm.script.tirx.select(condition, true_value, false_value)#

Construct a scalar conditional or select a host value.

Parameters:

condition (Expr or Python value) – Scalar IR condition or host truth value. Object-convertible conditions are normalized before selection.
true_value (Expr or Python value) – Value selected when the condition is true.
false_value (Expr or Python value) – Value selected when the condition is false.

Returns:

result – A primitive if_then_else expression, or the selected host object.

Return type:

Expr or Python value

Notes

Both Python value arguments are constructed before this call. The generated IR conditional evaluates only its selected arm at runtime.

The serial For statement.

Parameters:

start (Expr) – The minimum value of iteration.
stop (Expr) – The maximum value of iteration.
annotations (Dict[str, Any]) – The optional annotations of the For statement.
step (Expr) – The optional step value of iteration.
unroll (bool or int, optional) – If True, adds {"pragma_unroll": True} annotation, which asks CUDA codegen to emit #pragma unroll while preserving the loop as a C++ for. If False, adds {"disable_unroll": True} annotation. If a positive integer, emits #pragma unroll N. Boolean values are handled separately from integers, so False keeps disabling unrolling.
dtype (str, optional) – The dtype of the loop variable, either "int32" or "uint32". When omitted it is inferred from the bounds. Bounds that do not already have this dtype are converted (literals are retyped, other expressions get a Cast). Note T.thread_binding does not support this; its loop var is always int32.

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

tvm.script.tirx.set_mutable_cell_(target: TensorLoad | Var, value: Any, *, span: SpanEntry | Span | None = None) → AlreadyEmitted[Stmt]#

Implements tvm.script.ir_builder.parser_protocol.set_mutable_cell_().

Updates emit a scalar buffer store. Targets must denote scalar storage.

tvm.script.tirx.setattr_(target: Any, name: str, value: Any, *, span: SpanEntry | Span | None = None) → AlreadyEmitted[Stmt] | None#: Implements tvm.script.ir_builder.parser_protocol.setattr_().

tvm.script.tirx.setitem_(target: Any, key: Any, value: Any, *, span: SpanEntry | Span | None = None) → AlreadyEmitted[Stmt]#: Implements tvm.script.ir_builder.parser_protocol.setitem_().

tvm.script.tirx.shared_scalar(dtype: str = 'float32', *, annotations: dict[str, Any] | None = None) → TensorLoad#: Allocate a zero-dimensional buffer in shared memory.

tvm.script.tirx.shift_left(x, y, span=None)#

Return the result of x left shifted by y bits.

Parameters:

x (Expr) – Input argument.
y (Expr) – Input argument.

Returns:

z – The result.

Return type:

tvm.script.tirx.shift_right(x, y, span=None)#

Return the result of x right shifted by y bits.

Parameters:

x (Expr) – Input argument.
y (Expr) – Input argument.

Returns:

z – The result.

Return type:

tvm.script.tirx.sigmoid(x, *, ty=None, span=None)#

Quick function to get sigmoid

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.sin(x, *, ty=None, span=None)#

Take sin of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.sinh(x, *, ty=None, span=None)#

Take sinh of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.smem(shape: list[~tvm.ir.expr.Expr] | tuple[~tvm.ir.expr.Expr] | ~tvm.ir.expr.Expr | ~numbers.Integral, dtype: str = 'float32', data: ~tvm.ir.expr.Var | None = None, strides: list[~tvm.ir.expr.Expr] | None = None, elem_offset: ~tvm.ir.expr.Expr | None = None, byte_offset: ~tvm.ir.expr.Expr | None = None, *, scope: str = 'shared', align: int = -1, offset_factor: int = 0, layout: str | ~tvm.tirx.layout.Layout | None = <MISSING>, allocated_addr: int | tuple[int, ...] | None = None, annotations: dict[str, ~typing.Any] | None = None) → Var#

Allocate a tensor and return its variable.

Emits a Bind statement with a tirx.alloc_tensor Call:

buf = T.alloc_tensor((128, 128))

Parameters:

shape (Union[List[Expr], Tuple[Expr], Expr, Integral]) – The shape of the buffer to allocate.
dtype (str) – The data type of the buffer elements.
scope (str) – The storage scope of the buffer (e.g., “global”, “shared”).
data (Optional[Var]) – Optional explicit data pointer.
strides (Optional[List[Expr]]) – Optional strides.
elem_offset (Optional[Expr]) – Optional element offset.
byte_offset (Optional[Expr]) – Optional byte offset.
align (int) – Alignment requirement in bytes.
offset_factor (int) – Offset factor.
layout (Optional[Union[str, Layout]]) – Optional layout.
allocated_addr (Optional[Union[int, Tuple[int, ...]]]) – Optional pre-allocated address metadata.
annotations (Optional[Dict[str, Any]]) – Optional annotations for the allocation.

Returns:

res – The allocated buffer.

Return type:

tvm.script.tirx.sqrt(x, *, ty=None, span=None)#

Take square root of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.stack_alloca(dtype_str, num, *, ty=None, span=None)#

Return new on stack dtype[num]

Parameters:

dtype_str (str) – The data type of array.
num (int) – The size of array.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.stack_make_dltensor(data, shape, strides, ndim, arr_dtype, elem_offset, *, ty=None, span=None)#

Allocate a Tensor(DLTensor) on stack, return the handle

Parameters:

data (Expr) – The data of array.
shape (Expr) – The shape of array.
strides (Expr) – The strides of array.
ndim (Expr) – The dimensions of array.
arr_dtype (Expr) – The data type of array.
elem_offse (Expr) – The element offset of array.

Returns:

call – The call expression.

Return type:

tvm.script.tirx.stack_make_shape(*args, ty=None, span=None)#

Allocate a shape tuple on stack, return the handle

Parameters:: args (int) – The tuple shape.
Returns:: call – The call expression.
Return type:: Expr

tvm.script.tirx.tan(x, *, ty=None, span=None)#

Take tan of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.tanh(x, *, ty=None, span=None)#

Take hyperbolic tanh of input x.

Parameters:: x (Expr) – Input argument.
Returns:: y – The result.
Return type:: Expr

tvm.script.tirx.target(target_config: dict | str, host: dict | str | Target | None = None) → Target#

Create a target

Parameters:

target_config (Union[Dict, str]) – The target configuration.
host (Optional[Union[Dict, str, Target]]) – The target configuration.

Returns:

res – The target.

Return type:

Target

tvm.script.tirx.tensor_data_ptr(tensor, *, ty=None, span=None)#

Project a tensor variable’s physical pointer.

The result type is inferred from its element type and storage scope. ty may supply an explicit result type; span records the source location.

tvm.script.tirx.tensor_indices(buffer: Var, index)#

Translate logical flat or multidimensional indices for a concrete buffer.

A single index is unraveled in row-major logical order, retaining the outermost quotient. Explicit multidimensional coordinates pass through. The result indexes the original buffer, preserving its strides, layout, element offset and aliases; this function never creates a buffer view.

Parameters:

buffer (Var) – The concrete buffer whose logical shape determines the coordinates.
index (Expr or sequence of Expr) – A flat logical index or explicit multidimensional coordinates.

Returns:

indices – Coordinates for a buffer load or an explicitly emitted buffer store.

Return type:

list of Expr

tvm.script.tirx.tensor_store(dest: Expr, indices: list[Expr | slice], value: Expr) → AlreadyEmitted[Stmt]#

Emit a buffer store and return a receipt for the stored statement.

Parameters:

dest (Expr) – The destination expression.
indices (List[Union[Expr, slice]]) – The indices location to be stored.
value (Expr) – The value to be stored.

Returns:

result – Receipt for the exact stored statement; consuming it does not emit again.

Return type:

AlreadyEmitted[Stmt]

tvm.script.tirx.then_(*, span: SpanEntry | Span | None = None) → ThenFrame#: Implements tvm.script.ir_builder.parser_protocol.then_().

tvm.script.tirx.thread_binding(start: Expr, stop: Expr | None = None, thread: str | None = None, *, annotations: dict[str, Any] | None = None) → ForFrame#

The thread-binding For statement.

Parameters:

start (Expr) – The minimum value of iteration.
stop (Expr) – The maximum value of iteration.
thread (str) – The thread for loop variable to bind.
annotations (Dict[str, Any]) – The optional annotations of the For statement.

Returns:

res – The ForFrame.

Return type:

frame.ForFrame

tvm.script.tirx.thread_id(extents: list[Expr | int] | None = None, dtype: str = 'int32') → Var | tuple[Var, ...]#

Define a cta→thread scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

tvm.script.tirx.thread_id_in_wg(extents: list[Expr | int] | None = None, dtype: str = 'int32') → Var | tuple[Var, ...]#

Define a warpgroup→thread scope id. Pass None (the default) to defer the extent; it will be inferred at LowerTIRx from sibling closure.

dtype selects the dtype of the introduced vars ("int32" or "uint32").

tvm.script.tirx.throw_last_error(*, ty=None, span=None)#

Throw TVMGetLastError()

Returns:: ret – The return expression
Return type:: Expr

tvm.script.tirx.tmem(shape: list[~tvm.ir.expr.Expr] | tuple[~tvm.ir.expr.Expr] | ~tvm.ir.expr.Expr | ~numbers.Integral, dtype: str = 'float32', data: ~tvm.ir.expr.Var | None = None, strides: list[~tvm.ir.expr.Expr] | None = None, elem_offset: ~tvm.ir.expr.Expr | None = None, byte_offset: ~tvm.ir.expr.Expr | None = None, *, scope: str = 'tmem', align: int = -1, offset_factor: int = 0, layout: str | ~tvm.tirx.layout.Layout | None = <MISSING>, allocated_addr: int | tuple[int, ...] | None = None, annotations: dict[str, ~typing.Any] | None = None) → Var#

Allocate a tensor and return its variable.

Emits a Bind statement with a tirx.alloc_tensor Call:

buf = T.alloc_tensor((128, 128))

Parameters:

shape (Union[List[Expr], Tuple[Expr], Expr, Integral]) – The shape of the buffer to allocate.
dtype (str) – The data type of the buffer elements.
scope (str) – The storage scope of the buffer (e.g., “global”, “shared”).
data (Optional[Var]) – Optional explicit data pointer.
strides (Optional[List[Expr]]) – Optional strides.
elem_offset (Optional[Expr]) – Optional element offset.
byte_offset (Optional[Expr]) – Optional byte offset.
align (int) – Alignment requirement in bytes.
offset_factor (int) – Offset factor.
layout (Optional[Union[str, Layout]]) – Optional layout.
allocated_addr (Optional[Union[int, Tuple[int, ...]]]) – Optional pre-allocated address metadata.
annotations (Optional[Dict[str, Any]]) – Optional annotations for the allocation.

Returns:

res – The allocated buffer.

Return type:

tvm.script.tirx.trunc(*args, attrs=None, ty_args=None, span=None, ty=None, **kwargs)#: Construct a call to prim.trunc.

tvm.script.tirx.truncdiv(a, b, span=None)#

Compute the truncdiv of two expressions.

Parameters:

a (Expr) – The left hand operand
b (Expr) – The right hand operand
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type:

Note

This is the default integer division behavior in C.

tvm.script.tirx.truncmod(a, b, span=None)#

Compute the truncmod of two expressions.

Parameters:

a (Expr) – The left hand operand
b (Expr) – The right hand operand
span (Optional[Span]) – The location of this operator in the source.

Returns:

res – The result expression.

Return type: