tvm.relax.frontend

Frontends for constructing Relax programs, with the model importers

tvm.relax.frontend.detach_params(mod: IRModule) → Tuple[IRModule, Dict[str, List[NDArray]]]

Detach the attribute "params" in the functions of the input IRModule as separate dictionary of params.

参数:

mod (tvm.IRModule) -- The IRModule whose functions' "param" attribute is going to be detached.

返回:

• detached_mod (tvm.IRModule) -- The IRModule after the detachment.

• params_dict (Dict[str, List[tvm.nd.NDArray]]) -- The detached params. The dict keys corresponds to the names of the functions in the input IRModule that have attribute "params".

tvm.relax.frontend.nn

A PyTorch-like API to build IRModules.

class tvm.relax.frontend.nn.Any(*args, **kwargs)

Special type indicating an unconstrained type.

• Any is compatible with every type.

• Any assumed to have all methods.

• All values assumed to be instances of Any.

Note that all the above statements are true from the point of view of static type checkers. At runtime, Any should not be used with instance checks.

class tvm.relax.frontend.nn.Conv1D(in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, dtype: str | None = None)

Module for conv1d layer.

forward(x: Tensor) → Tensor

Forward method for conv1d layer.

参数:

x (Tensor) -- The input tensor.

返回:

ret -- The output tensor for the conv1d layer.

返回类型:

Tensor

class tvm.relax.frontend.nn.Conv2D(in_channels: int, out_channels: int, kernel_size: List[int] | int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, dtype: str | None = None, data_layout: str = 'NCHW')

Module for conv2d layer.

forward(x: Tensor) → Tensor

Forward method for conv2d layer.

参数:

x (Tensor) -- The input tensor.

返回:

ret -- The output tensor for the conv2d layer.

返回类型:

Tensor

class tvm.relax.frontend.nn.Conv3D(in_channels: int, out_channels: int, kernel_size: List[int] | int, stride: List[int] | int = 1, padding: List[int] | int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, dtype: str | None = None, data_layout: str = 'NCDHW')

Module for conv3d layer.

forward(x: Tensor) → Tensor

Forward method for conv3d layer.

参数:

x (Tensor) -- The input tensor.

返回:

ret -- The output tensor for the conv3d layer.

返回类型:

Tensor

class tvm.relax.frontend.nn.ConvTranspose1D(in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, output_padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, dtype: str | None = None)

Module for ConvTranspose1D layer.

forward(x: Tensor) → Tensor

Forward method for conv transpose 1d layer.

参数:

x (Tensor) -- The input tensor.

返回:

ret -- The output tensor for the conv transpose 1d layer.

返回类型:

Tensor

class tvm.relax.frontend.nn.Effect

Effect is a special non-user facing type that is used to represent operations with side effects, for example, print. It is used to represent the output of a computation.

create(name_hint: str) → List[Var]

Create the implicit inputs to a relax.Function that represents the side effect

emit_init(name_hint: str, builder: BlockBuilder) → List[DataflowVar]

Emit the initialization of the effect. This method is called by the compiler to initialize the effect.

finalize() → List[Var]

finalize the effect as the implicit return value of a relax.Function

set_state(state_vars: List[Var]) → None

Set the variables that represents the effect

to(dtype: str | None = None) → None

Convert the effect to specific dtype. Usually it is no-op for most of the effects

class tvm.relax.frontend.nn.Embedding(num: int | str | PrimExpr, dim: int | str | PrimExpr, dtype: str | None = None)

Module for embedding layer.

forward(x: Tensor)

Forward method for embedding layer.

参数:

x (Tensor) -- The input tensor.

返回:

ret -- The output tensor for the embedding layer.

返回类型:

Tensor

class tvm.relax.frontend.nn.ExternModule(symbols: Dict[str, Callable])

The abstract base class for external modules. External modules are designed to help incorporate user-provided handcrafted kernels into the exported TVM IRModule.

load() → Module

Loads the external module into a TVM runtime module.

class tvm.relax.frontend.nn.GELU

Module for GELU activation layer.

class tvm.relax.frontend.nn.GroupNorm(num_groups: int, num_channels: int, eps: float = 1e-05, affine: bool = True, dtype: str | None = None)

Module for group norm layer.

forward(x: Tensor, channel_axis: int = 1, axes: List[int] | None = None)

Forward method for group norm layer.

参数:

• x (Tensor) -- The input tensor.

• channel_axis (int) -- Channel axis of the input data.

• axes (Optional[List[int]]) -- Optional list of axes to compute norm over, if not specified, assumes that the first two axes should be left alone.

返回:

ret -- The output tensor for the group norm layer.

返回类型:

Tensor

class tvm.relax.frontend.nn.IOEffect

Modeling IO side effect, for example, printing the content of NDArrays on screen, inserting debug breakpoints, etc.

create(name_hint: str) → List[Var]

Create the implicit inputs to a relax.Function that represents the side effect

emit_init(name_hint, builder: BlockBuilder) → List[DataflowVar]

Emit the initialization of the effect. This method is called by the compiler to initialize the effect.

finalize() → List[Var]

finalize the effect as the implicit return value of a relax.Function

set_state(state_vars: List[Var]) → None

Set the variables that represents the effect

class tvm.relax.frontend.nn.KVCache(init_seq_len: int, unit_shape: Sequence[int], dtype: str | None = None)

Effect to implement KVCache.

append(new_element: Tensor) → None

Append a new element in KVCache.

参数:

new_element (Tensor) -- The new tensor to append.

create(name_hint: str) → List[Var]

Create the implicit inputs to a relax.Function that represents the KVCache effect.

参数:

name_hint (str) -- The name hint of the relax.Var.

返回:

ret -- The relax.Var for KVCache.

返回类型:

List[relax.Var]

emit_init(name_hint: str, bb: BlockBuilder)

Emit the initialization of the KVCache effect.

参数:

• name_hint (str) -- The name hint of the initialization binding Var.

• bb (relax.BlockBuilder) -- The relax BlockBuilder to emit.

finalize() → List[Var]

Finalize the KVCache effect as the implicit return value of a relax.Function.

返回:

ret -- The output relax.Var as KVCache.

返回类型:

List[rx.Var]

set_state(state_vars: List[Var]) → None

Set the variables that represents the effect

to(dtype: str | None = None) → None

Convert the KVCache effect to specific dtype.

参数:

dtype (Optional[str]) -- The target data type to convert.

view(seq_len: Var) → Tensor

View the last elements in KVCache.

参数:

seq_len (tir.Var) -- The number of last elements to view.

返回:

ret -- The last tensor to view.

返回类型:

Tensor

class tvm.relax.frontend.nn.LayerNorm(normalized_shape: int, eps: float | None = 1e-05, elementwise_affine: bool = True, dtype: str | None = None)

Module for Layer Normalization

forward(x: Tensor) → Tensor

Forward method for layer normalization layer.

参数:

x (Tensor) -- The input tensor.

返回:

ret -- The output tensor for the layer normalization layer.

返回类型:

Tensor

class tvm.relax.frontend.nn.Linear(in_features: int | str | PrimExpr, out_features: int | str | PrimExpr, bias: bool = True, dtype: str | None = None, out_dtype: str | None = None)

Module for linear layer.

forward(x: Tensor) → Tensor

Forward method for linear layer.

参数:

x (Tensor) -- The input tensor.

返回:

ret -- The output tensor for the linear layer.

返回类型:

Tensor

to(dtype: str | None = None) → None

Override to() such that we do not convert bias if there is out_dtype. Otherwise, we might run into dtype mismatch when computing x + self.bias since x is of type out_dtype and bias becomes dtype, potentially different.

class tvm.relax.frontend.nn.Module

Base class for neural network components. Subclass it to build your models. Modules can nest within each other in a tree structure using regular attribute assignment.

export_tvm(spec: _spec.ModuleSpecType, debug: bool = False, allow_extern: bool = False) → Tuple[IRModule, List[Tuple[str, Parameter]]] | Tuple[IRModule, List[Tuple[str, Parameter]], List[ExternModule]]

Export the module to TVM IRModule and parameters

参数:

• spec (_spec.ModuleSpecType) -- A dictionary mapping each input name to a specification that defines the inputs shape and dtype.

• debug (bool) -- If set to True, then the exported module will support effects. This enables things like printing in the graph.

返回:

• irmodule (tvm.ir.IRModule) -- The converted tvm IR representation of the model.

• params (List[Tuple[str, Parameter]]) -- A list of Parameters corresponding to the weights of the model.

• ext_mods (List[nn.ExternModule]) -- A list of ExternModules that are used in the model.

jit(spec: _spec.ModuleSpec, device: str | Device = 'cpu', pipeline: None | str | Pass = 'default_build', out_format: str = 'torch', debug: bool = False) → Any

Just-in-time compilation of a nn.model to an executable

load_state_dict(state_dict: Dict[str, Parameter], strict: bool = True) → Tuple[List[str], List[str]]

This function copies parameters and buffers from the state_dict into the current module and its descendants. If strict is set to True, the keys in the state_dict must exactly match the keys returned by the state_dict() function of this module.

参数:

• state_dict (Dict[str, Parameter]) -- A dictionary containing a whole state of the module

• strict (bool = True) -- Whether to strictly enforce that the keys in state_dict match the keys returned by this module's state_dict() function.

返回:

(missing_keys, unexpected_keys) -- A tuple of two lists: the missing keys and the unexpected keys.

返回类型:

Tuple[List[str], List[str]]

named_parameters(prefix: str = '') → Iterator[Tuple[str, Parameter]]

This method provides an iterator over module parameters, yielding both the parameter name and its corresponding value.

参数:

prefix (str) -- Prefix to prepend to all parameter names.

生成器:

(str, Parameter) - Tuple containing the name and parameter

parameters() → Iterator[Parameter]

This method provides an iterator over module parameters, yielding only the Parameter value.

生成器:

Parameter - The module's parameter

state_dict(*, prefix: str = '', destination: Dict[str, Parameter] | None = None) → Dict[str, Parameter]

Returns a dictionary containing references to the whole state of the module.

参数:

• prefix (str) -- Prefix to prepend to all parameter names.

• destination (Optional[Dict[str, Parameter]]) -- Dictionary to which state will be saved. If None, a new dictionary is created.

返回:

dict -- a dictionary containing a whole state of the module

返回类型:

Dict[str, Parameter]

to(dtype: str | None = None) → None

Convert the module to specific dtype recursively

class tvm.relax.frontend.nn.ModuleList(modules: List[Module])

Holds submodules in a list.

append(module: Module)

Add a module to the end of the ModuleList

forward(x)

Feed-forward pass of the module

to(dtype: str | None = None) → None

Convert the module to specific dtype recursively

class tvm.relax.frontend.nn.Mutator

The mutator for nn.Module transform. Users can override the visit_* methods to apply transform in different structures, or even override the visit method to change the logic of traversal.

visit(name: str, node: Any) → Any

The base dispatching method for visiting of all nodes.

参数:

• name (str) -- The name of the current node in parent's attribute.

• node (Any) -- The current node to visit.

返回:

ret_node -- The new node to replace current node.

返回类型:

Any

visit_effect(name: str, node: Parameter) → Any

The base visiting method for mutation of nn.Parameter nodes.

参数:

• name (str) -- The name of the current node in parent's attribute.

• node (nn.Parameter) -- The current node of nn.Parameter to mutate.

返回:

ret_node -- The new node to replace current node.

返回类型:

Any

visit_module(name: str, node: Module) → Any

The base visiting method for mutation of nn.Module nodes.

参数:

• name (str) -- The name of the current node in parent's attribute.

• node (nn.Module) -- The current node of nn.Module to mutate.

返回:

ret_node -- The new node to replace current node.

返回类型:

Any

visit_modulelist(name: str, node: ModuleList) → Any

The base visiting method for mutation of nn.ModuleList nodes.

参数:

• name (str) -- The name of the current node in parent's attribute.

• node (nn.ModuleList) -- The current node of nn.MoModuleListdule to mutate.

返回:

ret_node -- The new node to replace current node.

返回类型:

Any

visit_param(name: str, node: Effect) → Any

The base visiting method for mutation of nn.Effect nodes.

参数:

• name (str) -- The name of the current node in parent's attribute.

• node (nn.Effect) -- The current node of nn.Effect to mutate.

返回:

ret_node -- The new node to replace current node.

返回类型:

Any

class tvm.relax.frontend.nn.Object(*, _expr: RelayExpr, _name: str)

A wrapper on top of relax.Expr whose struct_info is the base ObjectStructInfo (rather than any its subclass). Object effectively represents non-tensor frontend components such as KV caches.

class tvm.relax.frontend.nn.ObjectModule(symbols: Dict[str, Callable], filepath: Path)

A subclass of nn.ExternModule, which allows users to provide an object .o file to be linked into compiled artifact;

load() → Module

Loads the external module into a TVM runtime module.

class tvm.relax.frontend.nn.Parameter(shape: Sequence[int | str | PrimExpr], dtype: str | None = None)

A parameter represents the weight of a neural network layer. It is a special tensor which could be bound or not bound to concrete values. If a parameter is bound to a concrete value, it is called a bound parameter, otherwise it is called an unbound parameter.

property data: NDArray | None

Returns the concrete value of the parameter if it is bound to a concrete value, otherwise returns None. The returned value is a tvm.runtime.NDArray.

to(dtype: str | None = None) → None

Change the dtype of the parameter if it is not bound to any concrete data

class tvm.relax.frontend.nn.RMSNorm(hidden_size: int, axes: int | List[int], epsilon: float = 1e-05, bias: bool = True, dtype: str | None = None)

Module for rms norm layer.

forward(x: Tensor)

Forward method for rms norm layer.

参数:

x (Tensor) -- The input tensor.

返回:

ret -- The output tensor for the rms norm layer.

返回类型:

Tensor

class tvm.relax.frontend.nn.ReLU

Module for ReLU activation layer.

class tvm.relax.frontend.nn.SiLU

Module for SiLU activation layer.

class tvm.relax.frontend.nn.SourceModule(symbols: Dict[str, Callable], source_code: str | Path, source_format: str, compile_options: List[str] | None = None, compiler: str | None = None, output_format: str = 'obj')

A subclass of nn.ExternModule. It compiles C++/CUDA source code and link them into the eventual IRModule.

Shape/dtype inference. The nn.ExternModule system requires users to provide additional information to work, namely, symbols. It is a dictionary that maps each symbol in the external object file to its shape/dtype inference function. Consider a case where function my_func accepts two tensors, a of shape (x, y, 1), and b of shape (y, z, 5), and produces a tensor c of shape (x, y, z, 9), the shape/dtype inference function should look like:

def shape_dtype_inference(a, b):
    x, y, _ = a.shape
    _, z, _ = b.shape
    return nn.Tensor.placeholder((x, y, z, 9), dtype="float32")

and the symbols dictionary should be provided as:

symbols={
    "my_func": shape_dtype_inference,
}

Calling convention. All external modules now follows "destination-passing-style" (DPS) calling convention, which means the returned tensors are pre-allocated by the system already and passed in as an argument of the external function.

Reuse the example above, the implementation of my_func should include three parameters in its signature, where tensors are represented using DLTensor from DLPack, the de facto standard of in-memory representation of tensors. More details: dmlc/dlpack.

To expose the symbol, TVM_DLL_EXPORT_TYPED_FUNC(symbol, function) is guaranteed available:

// those headers are guaranteed to be available
#include <dlpack/dlpack.h>
#include <tvm/runtime/data_type.h>
#include <tvm/runtime/packed_func.h>

namespace {
// anonymous namespace hides the symbol `_my_func_impl` from other translation units
int _my_func_impl(DLTensor* a, DLTensor* b, DLTensor* c) {
    // `a` and `b` are inputs, and `c` is the output
}
}
// expose symbol `my_func` instead of `_my_func_impl`
TVM_DLL_EXPORT_TYPED_FUNC(my_func, _my_func_impl);

A compiler pass `AttachExternModules`. It is introduced to attach a list of nn.ExternModule`s into an IRModule at any stage of the compilation pipeline, and attach the compiled external modules as `runtime.Module`s into IRModule's `external_mods attribute. It is required by linking in relax.build, but with the existence of this pass, source compilation can be deferred to arbitrary stage of TVM compilation.

Caveats. It is required to call nn.add_extern to register external modules exactly once during export_tvm. Each symbol should be registered exactly once to avoid potential conflicts, and otherwise an error will be raised.

compile(output_path: Path) → None

Compiles the source code in a provided directory and returns the compiled artifact.

static get_compile_options(source_format: str, tvm_pkg: List[str] | None = None) → List[str]

Returns the default compile options depending on source_format, including the default inlcude paths w.r.t. tvm_home(), default flags to configure DMLC-Core, and by default, it uses "-O3" and "-std=c++17".

参数:

• source_format (str) -- The source code format. It can be either "cpp" or "cu".

• tvm_pkg (Optional[List[str]]) -- The list of packages to be included under tvm_home/3rdparty. Each element should be a relative path to tvm_home/3rdparty.

返回:

compile_options -- The list of compilation flags.

返回类型:

List[str]

static get_includes(tvm_pkg: List[str] | None = None) → List[Path]

Returns the default include paths according to tvm_home(). By default, it includes TVM, DLPack, and DMLC-Core. With tvm_pkg provided, it also includes the specified package under tvm_home/3rdparty.

参数:

tvm_pkg (Optional[List[str]]) -- The list of packages to be included under tvm_home/3rdparty. Each element should be a relative path to tvm_home/3rdparty.

返回:

includes -- The list of include paths.

返回类型:

List[pathlib.Path]

load() → Module

Loads the external module into a TVM runtime module.

static tvm_home() → Path

Find TVM's home directory. If TVM_HOME environment variable is set, use it. Otherwise, use the directory where the tvm Python package is installed. As a sanity check, it is required to have include and 3rdparty as direct subdirectories.

返回:

tvm_home -- The TVM home directory, and it is guaranteed to have include and 3rdparty as direct subdirectories.

返回类型:

pathlib.Path

class tvm.relax.frontend.nn.SubroutineMixin

A mixin that generates a

Contains common logic for tvm.relax.frontend.nn.Module and tvm.relax.testing.nn.Module.

class tvm.relax.frontend.nn.Tensor(*, _expr: RelayExpr)

A wrapper on top of relax.Expr whose struct_info is a TensorStructInfo, providing more convenient access shape and dtype information. Tensor is always symbolc and not bound to any concrete values. Shape and dtype inference is done eagerly upon tensor creation, i.e. when operators are applied on tensors, the shape and dtype information is already available.

property dtype: str

Returns the data type of the tensor.

返回:

dtype -- The data type of the tensor

返回类型:

str

static from_const(data) → Tensor

Construct a tensor from numpy constants.

static from_scalar(data: int | float, dtype: str) → Tensor

Construct a tensor from a scalar with dtype specified.

static from_struct_info(struct_info: TensorStructInfo, name: str = 'tensor') → Tensor

Construct a nn.Tensor from relax TensorStructInfo

property ndim: int

Returns the number of dimensions of the tensor.

返回:

ndim -- The number of dimensions of the tensor

返回类型:

int

static placeholder(shape: Sequence[int | str | PrimExpr], dtype: str, name: str = 'tensor') → Tensor

Create a placeholder tensor with given shape and dtype. A placeholder tensor should never be created directly by users in usual cases, and the only exception is to indicate the shape/dtype of return values of an external function.

If shape is a string name, we create a symbolic shape tvm.tir.Var(name, "int64").

property shape: List[int | PrimExpr]

Returns the shape of the tensor as a list of integers.

An integer can be a python int or tvm.tir.PrimExpr, depending on whether the shape is fully static, for example, [1, 2, tvm.tir.Var("n")] is a valid shape where the last dimension is dynamic while the first two dimensions are always static constants.

返回:

shape -- The shape of the tensor

返回类型:

List[Union[int, tir.PrimExpr]]

class tvm.relax.frontend.nn.TypeVar

Type variable.

The preferred way to construct a type variable is via the dedicated syntax for generic functions, classes, and type aliases:

class Sequence[T]:  # T is a TypeVar
    ...

This syntax can also be used to create bound and constrained type variables:

# S is a TypeVar bound to str
class StrSequence[S: str]:
    ...

# A is a TypeVar constrained to str or bytes
class StrOrBytesSequence[A: (str, bytes)]:
    ...

However, if desired, reusable type variables can also be constructed manually, like so:

T = TypeVar('T')  # Can be anything
S = TypeVar('S', bound=str)  # Can be any subtype of str
A = TypeVar('A', str, bytes)  # Must be exactly str or bytes

Type variables exist primarily for the benefit of static type checkers. They serve as the parameters for generic types as well as for generic function and type alias definitions.

The variance of type variables is inferred by type checkers when they are created through the type parameter syntax and when infer_variance=True is passed. Manually created type variables may be explicitly marked covariant or contravariant by passing covariant=True or contravariant=True. By default, manually created type variables are invariant. See PEP 484 and PEP 695 for more details.

tvm.relax.frontend.nn.add(a: Tensor, b: Tensor, name: str = 'add') → Tensor

Addition with numpy-style broadcasting.

参数:

• a (Tensor) -- The first input tensor.

• b (Tensor) -- The second input tensor.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

示例

c = add(a, b)

tvm.relax.frontend.nn.add_extern(mod: ExternModule) → None

Add an external module to the exporter.

tvm.relax.frontend.nn.argsort(data: Tensor, axis: int = -1, descending: bool = False, dtype: str = 'int32', name='argsort')

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.

参数:

• data (Tensor) -- 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.

• name (str) -- Name hint.

返回:

out -- The indices of the sorted tensor.

返回类型:

Tensor

tvm.relax.frontend.nn.astype(x: Tensor, dtype: str, name: str = 'astype') → Tensor

Cast input tensor to the given data type.

参数:

• x (Tensor) -- The input data to the operator.

• dtype (str) -- The target data type

• name (str) -- Name hint.

返回:

result -- The casted result.

返回类型:

Tensor

tvm.relax.frontend.nn.broadcast_to(x: Tensor, shape: Sequence[int | PrimExpr], name: str = 'broadcast_to') → Tensor

Broadcasts a tensor to a specified shape.

参数:

• x (Tensor) -- The input data to the operator.

• shape (Sequence[IntExpr]) -- The target shape.

• name (str) -- Name hint.

返回:

result -- The broadcasted tensor.

返回类型:

Tensor

tvm.relax.frontend.nn.ccl_allgather(x: Tensor, num_workers: int, name='ccl_allgather')

CCL Allgather operator

参数:

• x (relax.Expr) -- The input tensor.

• num_workers (int) -- Number of workers.

• name (str) -- Name hint for this operation.

返回:

result -- The result tensor of allgather.

返回类型:

Tensor

tvm.relax.frontend.nn.ccl_allreduce(x: Tensor, op_type: str = 'sum', in_group: bool = True, name='ccl_allreduce')

CCL Allreduce operator

参数:

• x (relax.Expr) -- The input tensor.

• op_type (str) -- The type of reduction operation to be applied to the input data. Now "sum", "prod", "min", "max" and "avg" are supported.

• in_group (bool) -- Whether the reduction operation performs globally or in group as default.

• name (str) -- Name hint for this operation.

返回:

result -- The result tensor of allreduce.

返回类型:

Tensor

tvm.relax.frontend.nn.ccl_broadcast_from_worker0(x: Tensor, name='broadcast_from_worker')

Broadcast data from worker-0 to all other workers.

参数:

• x (Tensor) -- The tensor to be broadcast.

• name (str) -- Name hint for this operation.

返回:

result -- The same tensor, which has been broadcast to all other workers.

返回类型:

Tensor

tvm.relax.frontend.nn.chunk(x: Tensor, chunks: int, dim: int = 0, name: str = 'chunk') → Tensor

Split a tensor along dim into the specified number of chunks.

参数:

• x (Tensor) -- Input tensor to be split.

• chunks (int) -- Number of pieces to slice x into.

• dim (int) -- Which dimension to split x.

• name (str) -- Name hint for this operation.

返回:

result -- A tuple with chunks elements containing slices of x.

返回类型:

Tuple[Tensor]

tvm.relax.frontend.nn.concat(x: List[Tensor], dim: int, name: str = 'concat') → Tensor

Concatenate a list of tensors along an axis.

参数:

• x (List[Tensor]) -- List of tensors to concatenate.

• dim (int) -- Dimension to concatenate upon.

• name (str) -- Name hint for this operator.

返回:

result -- Expanded result.

返回类型:

Tensor

tvm.relax.frontend.nn.conv1d(x: Tensor, weight: Tensor, bias: Tensor | None = None, stride: int | Tuple | None = 1, padding: int | Tuple | str | None = 0, dilation: int | Tuple | None = 1, groups: int | None = 1, name: str = 'conv1d') → Tensor

1D convolution.

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

In the default case, where the data_layout is NCW and kernel_layout is OIW, conv1d takes in a data Tensor with shape (batch_size, in_channels, width), and a weight Tensor with shape (channels, in_channels, kernel_w), where kernel_w is the length of the W kernel dimension, to produce an output Tensor with the following rule:

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

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

参数:

• x (Tensor) -- The input data to the operator.

• weight (Tensor) -- The weight expressions.

• bias (Optional[Tensor]) -- Optional bias tensor of shape [O].

• strides (Optional[Union[int, Tuple]]) -- The strides of convolution. It is required to have length 1.

• padding (Optional[Union[int, Tuple, str]]) -- The padding of convolution on both sides of inputs before convolution. It is required to have length either 1 or 2.

• dilation (Optional[Union[int, Tuple]]) -- Specifies the dilation rate to be used for dilated convolution. It is required to have length 1.

• groups (Optional[int]) -- Number of groups to split the input into for grouped convolution. The number of input and output channels should be divisible by the number of groups.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.conv1d_transpose(x: Tensor, weight: Tensor, bias: Tensor | None = None, stride: int | Tuple[int] | None = 1, padding: int | Tuple[int, ...] | None = 0, output_padding: int | Tuple[int] | None = 0, dilation: int | Tuple | None = 1, groups: int | None = 1, name: str = 'conv1d_transpose') → Tensor

1D transposed convolution operator.

This operator can be seen as the gradient operator of conv1d.

The output shape can be explained in the simple case when data_layout == "NCW" and kernel_layout == "IOW". Suppose data has shape (N, in_channel, in_w), weight has shape (in_channel, out_channel, weight_w), we need to assure that in_channel % groups == 0. The shape of the output will be (N, out_channel * groups, out_w), where

• out_w = ((in_w - 1) * strides[0] + weight_w - 2 * padding[0] + output_padding[0])

参数:

• data (Tensor) -- The input data to the operator.

• weight (Tensor) -- The weight tensor.

• strides (Union[int, Tuple[int]]) -- The strides of convolution. It is required to have length 1.

• padding (Union[int, Tuple[int, ...]]) -- The padding of convolution on both sides of inputs before convolution. It is required to have length either 1 or 2.

• output_padding (Union[int, Tuple[int, ...]], optional) -- Used to disambiguate the output shape.

• dilation (Union[int, Tuple[int]]) -- Specifies the dilation rate to be used for dilated convolution. It is required to have length either 1.

• groups (int) -- Number of groups to split the input into for grouped convolution. The number of input and output channels should be divisible by the number of groups.

• data_layout (str) -- Layout of the input.

• kernel_layout (str) -- Layout of the weight.

• out_layout (Optional[str]) -- Layout of the output. If not specified, it is the same as data_layout

• out_dtype (Optional[Union[str, DataType]]) -- Specifies the output data type for mixed precision conv2d.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.conv2d(x: Tensor, weight: Tensor, bias: Tensor | None = None, stride: int | Tuple | None = 1, padding: int | Tuple | str | None = 0, dilation: int | Tuple | None = 1, groups: int | None = 1, data_layout: str | None = 'NCHW', name: str = 'conv2d') → Tensor

Applies a 2D convolution over an input image composed of sevaral input planes

参数:

• x (Tensor) -- Input tensor of shape [B, N, H, W]

• weight (Tensor) -- Filters of shape [O, N/groups, kH, kW]

• bias (Optional[Tensor]) -- Optional bias tensor of shape [O].

• stride (Optional[Union[int, Tuple]]) -- The stride of the convolving kernel. Can be a single number or tuple of (sH, sW).

• padding (Optional[[Union[int, Tuple]]]) -- Implicit paddings on both sides of the input.

• dilation (Optional[Union[int, Tuple]]) -- The spacing between kernel elements. Can be a single number of tuple (dH, dW).

• groups (Optional[int]) -- Split input into a number of groups.

• data_layout (Optional[str]) -- Layout of input and output data.

• name (str) -- Name hint.

返回:

result -- The computed result with shape [B, O, oH, oW].

返回类型:

Tensor

tvm.relax.frontend.nn.conv3d(x: Tensor, weight: Tensor, bias: Tensor | None = None, stride: int | Tuple | None = 1, padding: int | Tuple | str | None = 0, dilation: int | Tuple | None = 1, groups: int | None = 1, data_layout: str | None = 'NCDHW', name: str = 'conv3d') → Tensor

Applies a 3D convolution over an input image composed of sevaral input planes

参数:

• x (Tensor) -- Input tensor of shape [B, N, D, H, W]

• weight (Tensor) -- Filters of shape [O, N/groups, kD, kH, kW]

• bias (Optional[Tensor]) -- Optional bias tensor of shape [O].

• stride (Optional[Union[int, Tuple]]) -- The stride of the convolving kernel. Can be a single number or tuple of (sD, sH, sW).

• padding (Optional[[Union[int, Tuple]]]) -- Implicit paddings on both sides of the input.

• dilation (Optional[Union[int, Tuple]]) -- The spacing between kernel elements. Can be a single number of tuple (dD, dH, dW).

• groups (Optional[int]) -- Split input into a number of groups.

• data_layout (Optional[str]) -- Optional layout of the input and output data.

• name (str) -- Name hint.

返回:

result -- The computed result with shape [B, O, oD, oH, oW].

返回类型:

Tensor

tvm.relax.frontend.nn.cumsum(data: Tensor, axis: int | None = None, dtype: str | None = None, exclusive: bool | None = None, name: str = 'cumsum') → Tensor

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

参数:

• data (Tensor) -- 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[str]) -- 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 (Optional[bool]) -- If true will return exclusive sum in which the first element is not included.

• name (str) -- Name hint.

返回:

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.

返回类型:

Tensor

示例

a = [[1, 2, 3], [4, 5, 6]]

cumsum(a)  # if axis is not provided, cumsum is done over the flattened input.
-> [ 1,  3,  6, 10, 15, 21]

cumsum(a, dtype="float32")
-> [  1.,   3.,   6.,  10.,  15.,  21.]

cumsum(a, axis=0)  # sum over rows for each of the 3 columns
-> [[1, 2, 3],
    [5, 7, 9]]

cumsum(a, axis=1)
-> [[ 1,  3,  6],
    [ 4,  9, 15]]

a = [1, 0, 1, 0, 1, 1, 0]  # a is a boolean array
cumsum(a, dtype=int32)  # dtype should be provided to get the expected results
-> [1, 1, 2, 2, 3, 4, 4]

tvm.relax.frontend.nn.debug_func(name: str, *args: Tensor | PrimExpr | int | float | str, _line_info: str | None = None)

Call a debug function during runtime. The debug function must be registered with the following type signature:

@tvm.register_func(name_of_debug_func)
def debug_func(lineno: str, arg_0, arg_1, ...) -> None:
    ...

参数:

• name (str) -- The name of the debug function to call.

• *args (Union[Tensor, _tir.PrimExpr, int, float, str]) -- The arguments to pass to the debug function.

tvm.relax.frontend.nn.divide(a: Tensor, b: Tensor, name: str = 'divide') → Tensor

Division with numpy-style broadcasting.

参数:

• a (Tensor) -- The first input tensor.

• b (Tensor) -- The second input tensor.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

示例

c = divide(a, b)

tvm.relax.frontend.nn.empty(shape: Sequence[int | PrimExpr], dtype: str = 'float32', name: str = 'empty') → Tensor

Construct an uninitialized tensor, with the input shape and dtype.

参数:

• shape (Sequence[IntExpr]) -- The shape of the created tensor.

• dtype (str) -- The data type of the created tensor.

• name (str) -- Name hint.

返回:

result -- The result tensor.

返回类型:

Tensor

tvm.relax.frontend.nn.equal(a: Tensor, b: Tensor, name: str = 'equal') → Tensor

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

参数:

• a (Tensor) -- The first input tensor.

• b (Tensor) -- The second input tensor.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.exp(x: Tensor, name: str = 'exp') → Tensor

Applies the exponential function.

\[\text{Exp}(x) = e^x\]

参数:

• x (Tensor) -- The input data to the operator.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

备注

The input tensor is required to have float dtype

tvm.relax.frontend.nn.extern(name: str, args: Sequence[Tensor | PrimExpr | int | float | str], out: OutType) → OutType

Invoke an extern function during runtime. The extern function must be registered with the " TVM runtime using TVM_REGISTER_GLOBAL (C++), or tvm.register_func (Python).

参数:

• name (str) -- The name of the extern function to call.

• args (Sequence[Union[Tensor, _tir.PrimExpr, int, float, str]]) -- The arguments to pass to the extern function.

• out (Union[Tensor, List[Tensor]]) -- The output tensors, only

返回:

result -- The result

返回类型:

Tensor

tvm.relax.frontend.nn.full(shape: Sequence[int | PrimExpr], fill_value: Tensor, dtype: str = 'float32', name: str = 'full') → Tensor

Fill array with scalar value.

参数:

• shape (Sequence[IntExpr]) -- The shape of the created tensor.

• fill_value (Tensor) -- The value to fill. Must be a scalar tensor.

• dtype (str) -- The data type of the created tensor. If dtype is not given, it will by default use the dtype of fill_value.

• name (str) -- Name hint.

返回:

result -- The result tensor.

返回类型:

Tensor

tvm.relax.frontend.nn.gelu(x: Tensor, approximate: str | None = None, name: str = 'gelu') → Tensor

Applies the Gaussian Error Linear Units function

\[\text{GeLU}(x) = 0.5 * x * (1 + \text{erf}(x * 0.5**0.5))\]

where \(erf\) is the Gauss Error function.

参数:

• x (Tensor) -- The input data

• approximate (Optional[str]) -- If set to tanh, use an approximation when calculating CDF.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

备注

The input tensor is required to have float dtype

tvm.relax.frontend.nn.get_default_dtype() → str

Get the default parameter dtype if not specified. By default it is float32.

返回:

dtype -- The default dtype

返回类型:

str

tvm.relax.frontend.nn.get_timestep_embedding(x: Tensor, embedding_dim: int, flip_sin_to_cos: bool = False, downscale_freq_shift: float = 1, scale: float = 1, max_period: int = 10000, name: str = 'get_timestep_embedding') → Tensor

Timestep calculation as described in Denoising Diffusion Probabilistic Models.

参数:

• x (Tensor) -- A 1-D Tensor of N indices.

• embedding_dim (int) -- The dimension of the output.

• flip_sin_to_cos (bool) -- If True, change the order of sine and cosine embeddings.

• downscale_freq_shift (float) -- Adjusts the frequency of the sinusoidal sampling.

• scale (float) -- Weight adjustment for embedding magnitude.

• max_period (int) -- Controls the minimum frequency of the embeddings.

• name (str) -- The name to label this operator with.

返回:

result -- [N x dim] Tensor of positional embeddings.

返回类型:

Tensor

tvm.relax.frontend.nn.greater(a: Tensor, b: Tensor, name: str = 'greater') → Tensor

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

参数:

• a (Tensor) -- The first input tensor.

• b (Tensor) -- The second input tensor.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.greater_equal(a: Tensor, b: Tensor, name: str = 'greater_equal') → Tensor

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

参数:

• a (Tensor) -- The first input tensor.

• b (Tensor) -- The second input tensor.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.group_norm(x: Tensor, num_groups: int, weight: Tensor | None, bias: Tensor | None, eps: float = 1e-05, channel_axis: int = 1, axes: List[int] | None = None, name: str = 'group_norm') → Tensor

Applies Group Normalization over a mini-batch of inputs as described in the paper Group Normalization

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta\]

参数:

• x (Tensor) -- Input to which rms_norm will be applied.

• num_groups (int) -- Number of groups to separate the channels into.

• weight (Tensor) -- The gamma scale factor.

• bias (Tensor) -- The beta offset factor.

• epsilon (float) -- Small float added to square mean to avoid dividing by zero.

• channel_axis (int) -- The channel axis of the data.

• axes (Optional[int]) -- Which axes to compute the groupnorm over. If None, assumes first two channels should be ignored.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.interpolate(x: Tensor, size: int | Tuple[int] | None = None, scale_factor: float | Tuple[float] | None = None, mode: str = 'nearest', align_corners: bool | None = None, recompute_scale_factor: bool | None = None, antialias: bool | None = None, data_layout: str | None = 'NCHW', name: str = 'interpolate')

Resize a tensor using the specified mode.

参数:

• x (Tensor) -- Input tensor to be resized.

• size (Optional[Union[int, Tuple[int]]]) -- Requested output size, only one of size and scale_factor may be specified.

• scale_factor (Optional[Union[float, Tuple[float]]]) -- Multiplier for spatial size.

• mode (str) -- Algorithm used for sampling.

• align_corners (Optional[bool]) -- How to map pixels before and after sampling.

• recompute_scale_factor (Optional[bool]) -- Recompute the scale_factor for use in interpolation.

• antialias (Optional[bool]) -- Apply antialiasing to output.

• data_layout (Optional[str]) -- Layout of the input and output data.

• name (str) -- Name hint for this operation.

返回:

result -- Output tensor with requested shape.

返回类型:

Tensor

tvm.relax.frontend.nn.layer_norm(x: Tensor, normalized_shape: int | List[int], weight: Tensor | None = None, bias: Tensor | None = None, eps: float = 1e-05, name: str = 'layer_norm') → Tensor

Layer normalization (Lei Ba and et al., 2016). Applies layer normalization to the n-dimensional input array. This operator takes an n-dimensional input array and normalizes the input using the given axis:

\[out = \frac{data - mean(data, axis)}{\sqrt{var(data, axis)+\epsilon}} * gamma + beta\]

Unlike batch normalization, the mean and var are computed along the channel dimension.

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

备注

This operator can be optimized away for inference.

参数:

• x (Tensor) -- Input to which layer_norm will be applied.

• normalized_shape (Union[int, List[int]]) -- The shape of axes to normalize. If a single integer is used, it is treated as a singleton list and this module will normalize over the last dimension.

• weight (Tensor) -- The gamma scale factor.

• bias (Tensor) -- The beta offset factor.

• eps (float) -- Small float added to variance to avoid dividing by zero.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.less(a: Tensor, b: Tensor, name: str = 'less') → Tensor

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

参数:

• a (Tensor) -- The first input tensor.

• b (Tensor) -- The second input tensor.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.less_equal(a: Tensor, b: Tensor, name: str = 'less_equal') → Tensor

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

参数:

• a (Tensor) -- The first input tensor.

• b (Tensor) -- The second input tensor.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.matmul(a: Tensor, b: Tensor, out_dtype: str | None = None, name: str = 'matmul') → Tensor

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.

参数:

• a (Tensor) -- The first input tensor.

• b (Tensor) -- 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.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

示例

c = matmul(a, b)

tvm.relax.frontend.nn.maximum(x1: Tensor, x2: Tensor, name: str = 'maximum')

Element-wise maximum

参数:

• x1 (Tensor) -- The first input tensor.

• x2 (Tensor) -- The second input tensor.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

示例

c = maximum(a, b)

tvm.relax.frontend.nn.minimum(x1: Tensor, x2: Tensor, name: str = 'minimum')

Element-wise minimum

参数:

• x1 (Tensor) -- The first input tensor.

• x2 (Tensor) -- The second input tensor.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

示例

c = minimum(a, b)

tvm.relax.frontend.nn.multinomial_from_uniform(prob: Tensor, uniform_sample: Tensor, sample_indices: Tensor | None = None, dtype: str = 'int64', name: str = 'multinomial_from_uniform')

Returns a tensor where each row contains the index sampled from the multinomial probability distribution located in the corresponding row of tensor prob.

备注

For better cpu performance, use 'vm.builtin.multinomial_from_uniform'. For accurate results, ensure probabilities are between 0 and 1 and sum to 1.

参数:

• prob (Tensor) -- 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 (Tensor) -- The uniformly sampled 2-D tensor with the shape (n, 1). Values range from 0 to 1, indicating probabilities sampled uniformly.

• sample_indices (Optional[Tensor]) -- 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 output tensor.

返回:

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

返回类型:

Tensor

示例

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.relax.frontend.nn.multiply(a: Tensor, b: Tensor, name: str = 'mul') → Tensor

Multiplication with numpy-style broadcasting.

参数:

• a (Tensor) -- The first input tensor.

• b (Tensor) -- The second input tensor.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

示例

c = multiply(a, b)

tvm.relax.frontend.nn.negative(x: Tensor, name: str = 'neg') → Tensor

Numerical negative of the input tensor.

参数:

• x (Tensor) -- The input data to the operator.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.not_equal(a: Tensor, b: Tensor, name: str = 'not_equal') → Tensor

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

参数:

• a (Tensor) -- The first input tensor.

• b (Tensor) -- The second input tensor.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.ones(shape: Sequence[int | PrimExpr], dtype: str = 'float32', name: str = 'ones') → Tensor

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

参数:

• shape (Sequence[IntExpr]) -- The shape of the created tensor.

• dtype (str) -- The data type of the created tensor.

• name (str) -- Name hint.

返回:

result -- The result tensor.

返回类型:

Tensor

tvm.relax.frontend.nn.pad(x: Tensor, pad: List[int], mode: str = 'constant', value: int = 0, name: str = 'pad') → Tensor

Apply spatial padding to the input tensor.

参数:

• x (Tensor) -- Input tensor to be padded.

• pad (List[int]) -- List in the format of [before_0, after_0, before_1, after_1, ...] indicating how much to pad each axis of x.

• mod (str) -- Padding mode to use, constant implies padded elements will use value argument.

• value (int) -- What to pad with in constant mode.

• name (str) -- Name hint for this operator.

返回:

result -- Padded output tensor.

返回类型:

Tensor

tvm.relax.frontend.nn.permute(x: Tensor, axes: List[int] | None, name: str = 'permute') → Tensor

Permutes the dimensions of the input tensor.

参数:

• x (Tensor) -- The input data to the operator.

• axes (Optional[List[int]]) -- The target axes order.

• name (str) -- Name hint.

返回:

result -- The transposed result.

返回类型:

Tensor

tvm.relax.frontend.nn.permute_dims(x: Tensor, axes: List[int] | None = None, name: str = None) → Tensor

Permutes the dimensions of an array.

参数:

• x (Tensor) -- The input data to the operator.

• axes (Optional[List[int]]) -- The target axes order, reverse order if not specified.

• name (str) -- Name hint.

返回:

result -- The transposed result.

返回类型:

Tensor

tvm.relax.frontend.nn.print_(tensor: Tensor)

Debug printing a Tensor during runtime.

tvm.relax.frontend.nn.relu(x: Tensor, name: str = 'relu') → Tensor

Rectified Linear Unit (ReLU) activation function.

\[ext{ReLU}(x) = ext{max}(x, 0)\]

参数:

• x (Tensor) -- The input data.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.renormalize_top_p_top_k_prob(prob, sorted_prob, top_p, top_k)

Renormalizes probabilities after filtering with top_p and top_k, ensuring they sum up to 1.

备注

For accurate results, ensure probabilities are between 0 and 1 and sum to 1.

参数:

• prob (Tensor) -- A 2-D tensor of shape (batch, vocab_size) representing probability distributions.

• sorted_prob (Tensor) -- Probabilities sorted in descending order.

• top_p (Tensor) -- The cumulative probability threshold with shape (batch, 1) for nucleus sampling.

• top_k (Tensor) -- A tensor with shape (batch, 1), representing the number of top probabilities to consider for top-k sampling.

返回:

result -- The filtered and nomalized tensor with the sampe shape as input prob.

返回类型:

Tensor

tvm.relax.frontend.nn.repeat(x: Tensor, repeats: int, axis: int | None = None, name='repeat') → Tensor

Repeats elements of an array.

参数:

• data (Tensor) -- 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.

• name (str) -- Name hint.

返回:

ret -- The computed result.

返回类型:

Tensor

示例

np_x = numpy.array([[1, 2], [3, 4]])
x = Tensor.from_const(np_x)
lv1 = repeat(x, repeats=2) # lv1 == [1, 1, 2, 2, 3, 3, 4, 4]
lv2 = repeat(x, repeats=2, axis=1)   # lv2 == [[1., 1., 2., 2.],
                                     #         [3., 3., 4., 4.]]

tvm.relax.frontend.nn.reshape(x: Tensor, shape: Sequence[int | PrimExpr], name='reshape') → Tensor

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,)

参数:

• x (Tensor) -- The input data to the operator.

• shape (Sequence[IntExpr]) -- The new shape. Should be compatible with the original shape.

• name (str) -- Name hint.

返回:

result -- The reshaped result.

返回类型:

Tensor

备注

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.relax.frontend.nn.rms_norm(x: Tensor, weight: Tensor, axes: int | List[int], epsilon: float = 1e-05, name: str = 'rms_norm') → Tensor

Root mean square normalization (Biao Zhang and et al., 2019). Applies root mean square normalization to the n-dimensional input array. This operator takes an n-dimensional input array and normalizes the input using the given axis:

\[out = \frac{data}{\sqrt{mean(data, axis)+\epsilon}} * weight\]

参数:

• data (Tensor) -- Input to which rms_norm will be applied.

• weight (Tensor) -- The scale factor.

• axes (Union[int, List[int]]) -- The axes that along which the normalization is applied.

• epsilon (float) -- Small float added to square mean to avoid dividing by zero.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.sample_top_p_top_k_from_sorted_prob(sorted_prob: Tensor, sorted_index: Tensor, top_p: Tensor, top_k: Tensor, uniform_sample: Tensor, sample_indices: Tensor | None = None)

Samples indices from a sorted probability tensor based on top_p and top_k criteria.

备注

For accurate results, ensure probabilities are between 0 and 1 and sum to 1.

参数:

• sorted_prob (Tensor) -- A 2-D tensor, with shape (batch, vocab_size), contains probabilities sorted in descending order.

• sorted_index (Tensor) -- The indices tensor with shape (batch, vocab_size), corresponding to the sorted_prob. Potentially from applying argsort on the original probability tensor in descending order.

• top_p (Tensor) -- The cumulative probability threshold with shape (batch, 1) for nucleus sampling.

• top_k (Tensor) -- A tensor with shape (batch, 1), representing the number of top probabilities to consider for top-k sampling.

• uniform_sample (Tensor) -- Uniformly sampled values with shape (n, 1) are used to select the output indices.

• sample_indices (Optional[Tensor]) -- 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].

返回:

result -- The selected indices with shape (n, 1).

返回类型:

Tensor

示例

prob = [[0.1 , 0.4, 0.5],
        [0.3, 0.3, 0.4]]
sorted_prob = [[0.5, 0.4, 0.1],
               [0.4, 0.3, 0.3]]
sorted_index = [[2, 1, 0],
                [2, 0, 1]]
top_p = [[0.6],[0.9]]
top_k = [[3],[2]]
uniform_sample = [[0.5], [0.6]]
sample_indices = [[0], [1]]

sample_top_p_top_k_from_sorted_prob(
    sorted_prob, sorted_index,top_p, top_k, uniform_sample, sample_indices)
-> [2, 0]

tvm.relax.frontend.nn.scaled_dot_product_attention(query: Tensor, key: Tensor, value: Tensor, attn_mask: Tensor | None = None, is_causal: bool | None = False, scale: float | None = None, name: str = 'scaled_dot_product_attention')

Computes a scaled dot product attention on provided attention query, key, and values. Compliant with the functional torch implementation.

参数:

• query (Tensor) -- Tensor representing current attention lookup of shape [batch, seq_len, num_heads, head_size].

• key (Tensor) -- Tensor representing cross attention mapping of shape [batch, seq_len_kv, num_heads_kv, head_size].

• value (Tensor) -- Tensor representing embedded attention values of shape [batch, seq_len_kv, num_heads_kv, head_size_value].

• attn_mask (Optional[Tensor]) -- Optional mask for attention, not yet supported.

• is_causal (Optional[bool]) -- If set, uses a causal attention mask.

• scale (Optional[float]) -- Optional extra scaling argument applied to attention.

• name (str) -- Name hint for this function.

tvm.relax.frontend.nn.sigmoid(x: Tensor, name: str = 'sigmoid') → Tensor

Computes sigmoid.

\[\text{sigmoid}(x) = \frac{1}{1 + \exp(-x)}\]

参数:

• data (Tensor) -- The input data to the operator.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

备注

The input tensor is required to have float dtype

tvm.relax.frontend.nn.silu(x: Tensor, name: str = 'silu') → Tensor

Sigmoid Linear Unit function

\[\text{SiLU}(x) = x * \text{sigmoid}(x)\]

参数:

• data (Tensor) -- The input data

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

备注

The input tensor is required to have float dtype

tvm.relax.frontend.nn.softmax(x: Tensor, axis: int = -1, name: str = 'softmax') → Tensor

Computes softmax.

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

参数:

• data (Tensor) -- The input data to the operator.

• axis (int) -- The axis to sum over when computing softmax. If not specified, it is by default the last axis of the input tensor. Supports negative indexing.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

备注

The input tensor is required to have float dtype

tvm.relax.frontend.nn.sort(x: Tensor, axis: int = -1, descending: bool = False, name='sort')

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

参数:

• x (Tensor) -- 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

• name (str) -- Name hint.

返回:

out -- The sorted tensor.

返回类型:

Tensor

tvm.relax.frontend.nn.split(ary: Tensor, indices_or_sections: int | Sequence[int], axis: int = 0, name: str = 'split') → Tuple[Tensor, ...]

Split an array into multiple sub-arrays.

参数:

• ary (Tensor) -- Input tensor to be split.

• indices_or_sections (Union[int, Sequence[int]]) -- Indices or sections to split into.

• axis (int = 0) -- The axis along which to split, default is 0.

• name (str) -- Name hint.

返回:

result -- A list of sub-arrays as the outcome of splitting.

返回类型:

Tuple[Tensor, ...]

tvm.relax.frontend.nn.sqrt(x: Tensor, name: str = 'sqrt') → Tensor

Computes the element-wise sqrt of the input tensor.

参数:

• x (Tensor) -- The input tensor.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

备注

The input tensor is required to have float dtype

tvm.relax.frontend.nn.square(x: Tensor, name: str = 'square') → Tensor

Computes the element-wise square of the input tensor.

参数:

• x (Tensor) -- The input tensor.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.squeeze(x: Tensor, axis: int = -1, name: str = 'squeeze') → Tensor

Squeeze axes in the array.

参数:

• x (Tensor) -- 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.

• name (str) -- Name hint.

返回:

result -- The squeezed result.

返回类型:

Tensor

tvm.relax.frontend.nn.subtract(a: Tensor, b: Tensor, name: str = 'subtract') → Tensor

Subtraction with numpy-style broadcasting.

参数:

• a (Tensor) -- The first input tensor.

• b (Tensor) -- The second input tensor.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

示例

c = subtract(a, b)

tvm.relax.frontend.nn.sum(x: Tensor, axis: int | List[int] | None = None, keepdims: bool = False, name: str = 'sum') → Tensor

Computes the sum of tensor elements over given axes.

参数:

• x (Tensor) -- 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.

• name (str) -- Name hint for this operation.

返回:

result -- The computed result.

返回类型:

Tensor

tvm.relax.frontend.nn.take(x: Tensor, indices: Tensor, axis: int | None = None, name='take') → Tensor

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).

参数:

• x (Tensor) -- The source tensor.

• indices (Tensor) -- 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.

• name (str) -- Name hint.

返回:

ret -- The taken result.

返回类型:

Tensor

tvm.relax.frontend.nn.tanh(x: Tensor, name: str = 'tanh') → Tensor

Applies the hyperbolic tangent function.

\[\text{Tanh}(x) = \frac{e^x - e^{-x}}{e^x + e^{-x}}\]

参数:

• x (Tensor) -- The input data to the operator.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Tensor

备注

The input tensor is required to have float dtype

tvm.relax.frontend.nn.tensor_expr_op(tensor_expr_func: Callable, name_hint: str, args: List[Tensor | Var | int], *, attrs: Dict[str, Any] | None = None)

Build the given tensor_expr_func with te.

参数:

• tensor_expr_func (Callable) -- A function that returns a te tensor or a list of tensors.

• name_hint (str) -- Name hint.

• args (List[Union[Tensor, _tir.Var]]) -- Arguments passed to the function.

• attrs (Optional[Dict[str, Any]]) -- A dict of attributes to apply to the function.

返回:

result -- The result tensor.

返回类型:

Tensor

tvm.relax.frontend.nn.tensor_ir_inplace_op(func: PrimFunc, name_hint: str, args: Tensor | Sequence[Tensor | ShapeExpr | PrimExpr], inplace_indices: int | List[int], out: OutType) → OutType

Create a call_tir_inplace binding with given PrimFunc

参数:

• func (_tir.PrimFunc) -- The PrimFunc to call.

• name_hint (str) -- Name hint.

• args (Union[Tensor, Sequence[Union[Tensor, rx.ShapeExpr, _tir.PrimExpr]]]) -- The arguments to pass to the PrimFunc.

• 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.

• out (Union[Tensor, List[Tensor]]) -- The output tensors.

返回:

result -- The result tensor

返回类型:

Tensor

tvm.relax.frontend.nn.tensor_ir_op(func: PrimFunc, name_hint: str, args: Tensor | Sequence[Tensor | ShapeExpr | PrimExpr], out: OutType) → OutType

Create a call_tir binding with given PrimFunc

参数:

• func (_tir.PrimFunc) -- The PrimFunc to call.

• name_hint (str) -- Name hint.

• args (Union[Tensor, Sequence[Union[Tensor, rx.ShapeExpr, _tir.PrimExpr]]]) -- The arguments to pass to the PrimFunc.

• out (Union[Tensor, List[Tensor]]) -- The output tensors.

返回:

result -- The result tensor

返回类型:

Tensor

tvm.relax.frontend.nn.topk(data: Tensor, k: int = 1, axis: int = -1, ret_type: str = 'both', largest: bool = True, dtype: str = 'int32', name: str = 'topk')

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").

参数:

• data (Tensor) -- 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.

• name (str) -- Name hint.

返回:

out -- The computed result.

返回类型:

Tensor or Tuple[Tensor, Tensor]

tvm.relax.frontend.nn.triu(x: Tensor, diagonal: int = 0, name: str = 'triu') → Tensor

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

参数:

• x (Tensor) -- 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.

• name (str) -- Name hint.

返回:

ret -- The result tensor.

返回类型:

Tensor

tvm.relax.frontend.nn.unsqueeze(x: Tensor, dim: int, name: str = 'unsqueeze') → Tensor

Add a new axis to a tensor

参数:

• x (Tensor) -- Input tensor to expand.

• dim (int) -- Dimension to expand.

• name (str) -- Name hint for this operator.

返回:

result -- Expanded result.

返回类型:

Tensor

tvm.relax.frontend.nn.where(condition: Tensor, x1: Tensor, x2: Tensor, name: str = 'where') → Tensor

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.

参数:

• condition (Tensor) -- When True, yield x1; otherwise, yield x2. Must be broadcasting compatible with x1 and x2. Must have boolean dtype.

• x1 (Tensor) -- The first input tensor. Must be broadcasting compatible with condition and x2.

• x2 (Tensor) -- The second input tensor. Must be broadcasting compatible with condition and x1.

• name (str) -- Name hint.

返回:

result -- The result tensor.

返回类型:

Tensor

tvm.relax.frontend.nn.wrap_nested(expr: RelayExpr, name: str) → Tensor | Sequence[Tensor]

Wrap the given relax.Expr, emit it using the current BlockBuilder, and automatically handle nested cases if the expr represents a Tuple.

参数:

• expr (relax.Expr) -- The Expr to be wrapped.

• name (str) -- Name hint.

返回:

result -- The computed result.

返回类型:

Union[Tensor, Tuple[Tensor]]

tvm.relax.frontend.nn.zeros(shape: Sequence[int | PrimExpr], dtype: str = 'float32', name: str = 'zeros') → Tensor

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

参数:

• shape (Sequence[IntExpr]) -- The shape of the created tensor.

• dtype (str) -- The data type of the created tensor.

• name (str) -- Name hint.

返回:

result -- The result tensor.

返回类型:

Tensor

tvm.relax.frontend.onnx

Tools for converting ONNX graphs into Relax graphs.

tvm.relax.frontend.onnx.from_onnx(model: GraphProto, shape_dict: Dict[str, List] | None = None, dtype_dict: str | Dict[str, str] | None = 'float32', opset: int = None, keep_params_in_input: bool = False, sanitize_input_names: bool = True) → IRModule

Convert a ONNX model into an equivalent Relax Function. ONNX graphs are represented as Python Protobuf objects.

The current implementation assumes that the input model is after ONNX v1.1.0.

参数:

• model (protobuf object) -- ONNX ModelProto after ONNX v1.1.0

• shape_dict (dict of str to tuple, optional) -- The input shape to the graph

• dtype_dict (str or dict of str to str, optional) -- The input types to the graph

• opset (int, optional) -- Override to autodetected opset. This can be helpful for some testing.

• keep_params_in_input (bool) -- If True, parameters will be treated as input variables. If false, parameters are treated as constant and folded directly into the graph.

• sanitize_input_names (bool, optional) -- Whether to sanitize the input names to ensure they are valid Relax identifiers.

返回:

mod -- The relax module for compilation

返回类型:

tvm.IRModule

tvm.relax.frontend.stablehlo

StableHLO Frontends for constructing Relax programs, with the model importers

tvm.relax.frontend.stablehlo.from_stablehlo(stablehlo_module, input_info: List[Tuple[Tuple[int], str]] = None) → IRModule

Convert a StableHLO Module to a Relax program

参数:

• stablehlo_module (Union[str, mlir.ir.Module]) -- The StableHLO Module to convert.

• input_info (List[Tuple[Tuple[int], str]]) -- A list of shapes and data types of input tensors.

返回:

output -- The result IRModule with entry function "main"

返回类型:

tvm.IRModule

tvm.relax.frontend.torch