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.

Parameters:

mod (tvm.IRModule) – The IRModule whose functions’ “param” attribute is going to be detached.

Returns:

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

relax.frontend.nn.Module for conv1d layer.

forward(x: Tensor) → Tensor

Forward method for conv1d layer.

Parameters:: x (Tensor) – The input tensor.
Returns:: ret – The output tensor for the conv1d layer.
Return type:: 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)

relax.frontend.nn.Module for ConvTranspose1D layer.

forward(x: Tensor) → Tensor

Forward method for conv transpose 1d layer.

Parameters:: x (Tensor) – The input tensor.
Returns:: ret – The output tensor for the conv transpose 1d layer.
Return type:: 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

relax.frontend.nn.Module for embedding layer.

forward(x: Tensor)

Forward method for embedding layer.

Parameters:: x (Tensor) – The input tensor.
Returns:: ret – The output tensor for the embedding layer.
Return type:: 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: relax.frontend.nn.Module for GELU activation layer.

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.

Parameters:: 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.

Parameters:: name_hint (str) – The name hint of the relax.Var.
Returns:: ret – The relax.Var for KVCache.
Return type:: List[relax.Var]

emit_init(name_hint: str, bb: BlockBuilder)

Emit the initialization of the KVCache effect.

Parameters:

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.

Returns:: ret – The output relax.Var as KVCache.
Return type:: 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.

Parameters:: dtype (Optional[str]) – The target data type to convert.

view(seq_len: Var) → Tensor

View the last elements in KVCache.

Parameters:: seq_len (tir.Var) – The number of last elements to view.
Returns:: ret – The last tensor to view.
Return type:: Tensor

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

relax.frontend.nn.Module for Layer Normalization

forward(x: Tensor) → Tensor

Forward method for layer normalization layer.

Parameters:: x (Tensor) – The input tensor.
Returns:: ret – The output tensor for the layer normalization layer.
Return type:: Tensor

relax.frontend.nn.Module for linear layer.

forward(x: Tensor) → Tensor

Forward method for linear layer.

Parameters:: x (Tensor) – The input tensor.
Returns:: ret – The output tensor for the linear layer.
Return type:: 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

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.

Returns:

irmodule (tvm.ir.IRModule) – The converted tvm IR representation of the model.
params (Dict[str, tvm.nd.array]) – A dictionary of parameters corresponding to the weights of the model.
ext_mods (List[nn.ExternModule])

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.

Parameters:

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.

Returns:

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

Return type:

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.

Parameters:: prefix (str) – Prefix to prepend to all parameter names.
Yields:: (str, Parameter) - Tuple containing the name and parameter

parameters() → Iterator[Parameter]

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

Yields:: 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.

Parameters:

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.

Returns:

dict – a dictionary containing a whole state of the module

Return type:

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

Parameters:

name (str) – The name of the current node in parent’s attribute.
node (Any) – The current node to visit.

Returns:

ret_node – The new node to replace current node.

Return type:

Any

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

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

Parameters:

name (str) – The name of the current node in parent’s attribute.
node (nn.Parameter) – The current node of nn.Parameter to mutate.

Returns:

ret_node – The new node to replace current node.

Return type:

Any

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

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

Parameters:

name (str) – The name of the current node in parent’s attribute.
node (nn.Module) – The current node of nn.Module to mutate.

Returns:

ret_node – The new node to replace current node.

Return type:

Any

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

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

Parameters:

name (str) – The name of the current node in parent’s attribute.
node (nn.ModuleList) – The current node of nn.MoModuleListdule to mutate.

Returns:

ret_node – The new node to replace current node.

Return type:

Any

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

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

Parameters:

name (str) – The name of the current node in parent’s attribute.
node (nn.Effect) – The current node of nn.Effect to mutate.

Returns:

ret_node – The new node to replace current node.

Return type:

Any

class tvm.relax.frontend.nn.Object(*, _expr: Expr, _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)

relax.frontend.nn.Module for rms norm layer.

forward(x: Tensor)

Forward method for rms norm layer.

Parameters:: x (Tensor) – The input tensor.
Returns:: ret – The output tensor for the rms norm layer.
Return type:: Tensor

class tvm.relax.frontend.nn.ReLU: relax.frontend.nn.Module for ReLU activation layer.

class tvm.relax.frontend.nn.SiLU: relax.frontend.nn.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: https://github.com/dmlc/dlpack/blob/v0.8/include/dlpack/dlpack.h#L163-L206.

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

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

Parameters:

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.

Returns:

compile_options – The list of compilation flags.

Return type:

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.

Parameters:: 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.
Returns:: includes – The list of include paths.
Return type:: 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.

Returns:: tvm_home – The TVM home directory, and it is guaranteed to have include and 3rdparty as direct subdirectories.
Return type:: 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: Expr)

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.

Returns:: dtype – The data type of the tensor
Return type:: 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.

property ndim: int

Returns the number of dimensions of the tensor.

Returns:: ndim – The number of dimensions of the tensor
Return type:: 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.

Returns:: shape – The shape of the tensor
Return type:: List[Union[int, tir.PrimExpr]]

class tvm.relax.frontend.nn.TypeVar(name, *constraints, bound=None, covariant=False, contravariant=False)

Type variable.

Usage:

T = TypeVar('T')  # Can be anything
A = TypeVar('A', str, bytes)  # Must be 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 definitions. See class Generic for more information on generic types. Generic functions work as follows:

def repeat(x: T, n: int) -> List[T]:
‘’’Return a list containing n references to x.’’’ return [x]*n

def longest(x: A, y: A) -> A:
‘’’Return the longest of two strings.’’’ return x if len(x) >= len(y) else y

The latter example’s signature is essentially the overloading of (str, str) -> str and (bytes, bytes) -> bytes. Also note that if the arguments are instances of some subclass of str, the return type is still plain str.

At runtime, isinstance(x, T) and issubclass(C, T) will raise TypeError.

Type variables defined with covariant=True or contravariant=True can be used to declare covariant or contravariant generic types. See PEP 484 for more details. By default generic types are invariant in all type variables.

Type variables can be introspected. e.g.:

T.__name__ == ‘T’ T.__constraints__ == () T.__covariant__ == False T.__contravariant__ = False A.__constraints__ == (str, bytes)

Note that only type variables defined in global scope can be pickled.

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

Addition with numpy-style broadcasting.

Parameters:

a (Tensor) – The first input tensor.
b (Tensor) – The second input tensor.
name (str) – Name hint.

Returns:

result – The computed result.

Return type:

tvm.relax.frontend¶

tvm.relax.frontend.nn¶

tvm.relax.frontend.onnx¶

tvm.relax.frontend.stablehlo¶

tvm.relax.frontend.torch¶