Skip to main content

module

Module primitives.

A Module is a very basic container that holds Graphs for the purpose of compiling or exporting them as a whole.

Note: Modules are not to be confused with layers or modules found in other high-level APIs (like torch.nn.Module, or tf.Module). The MAX Graph API is a low level library. Rather, a Module is closer to ONNX Models.

Module

A Module is a container that holds Graphs.

Implemented traits:

AnyType, Copyable, Movable, Stringable

Methods:

__init__

__init__(inout self: Self)

Constructs an empty Module instance.

__str__

__str__(self: Self) -> String

Constructs a human-readable string representation of the module.

The string is in MLIR text format, and will include representations of any graphs and ops inside it.

Returns:

A human-readable string representation of the module.

verify

verify(self: Self)

Verifies the Module and its contents.

Examples of cases when a Graph may not be valid (the list is not exhaustive):

  1. it has an output op whose types don't match its out_types
  2. it has an op with an invalid name, number, type of operands, output types, etc.
  3. it contains cycles

Raises: If the Module did not pass verification. In this case it will also print a diagnostic message indicating the error.

save_to_file

save_to_file(self: Self, path: Path)

Saves this Module to a file.

Important: The file contents uses internal MLIR Bytecode format that is not guaranteed to be cross-version compatible and may change without notice.

Args:

  • path (Path): The path to save the Module to.

tensor_attr

tensor_attr[dtype: DType](self: Self, name: String, owned value: Tensor[dtype]) -> NamedAttribute

Creates a new Tensor-valued Attribute.

The value of this attribute will have the type MOTensor with the same shape and dtype as value. This method takes ownership of value and is suitable for use with very large Tensor values (such as model weights).

Parameters:

  • dtype (DType): The attribute tensor's element type.

Args:

  • name (String): The Attribute name.
  • value (Tensor[dtype]): The Attribute value.

Returns:

An internal representation of an Attribute.

tensor_resource_attr

tensor_resource_attr(self: Self, name: String, file_name: String, type: MOTensor) -> NamedAttribute

Creates a new Tensor Attribute from an external file.

The value of this constant will have the type type. The file must contain the Tensors raw data, as returned by Tensor.data. No endianness transformation is performed.

Args:

  • name (String): The Attribute name.
  • file_name (String): The file name to load from.
  • type (MOTensor): The Tensor type (element type, shape).

Returns:

An internal representation of an Attribute.

vector_attr

vector_attr[dtype: DType](self: Self, name: String, values: List[SIMD[dtype, 1]]) -> NamedAttribute

Creates a new Tensor-valued Attribute.

The value of this attribute will have the type MOTensor with 1D shape, consistent with the size of values.

Parameters:

  • dtype (DType): The attribute tensor's element type.

Args:

  • name (String): The Attribute name.
  • values (List[SIMD[dtype, 1]]): A vector representing the attribute's value.

Returns:

An internal representation of an Attribute.

scalar_attr

scalar_attr[dtype: DType](self: Self, name: String, value: SIMD[dtype, 1], rank: Int) -> NamedAttribute

Creates a new Tensor-valued Attribute.

The Tensor is considered to contain a single element, and its shape be of the specified rank (for example, rank=0 denotes a scalar).

Parameters:

  • dtype (DType): The attribute tensor's element type.

Args:

  • name (String): The Attribute name.
  • value (SIMD[dtype, 1]): The Attribute value.
  • rank (Int): The attribute tensor's rank.

Returns:

An internal representation of an Attribute.

string_attr

string_attr(self: Self, name: String, value: String) -> NamedAttribute

Creates a new String-valued Attribute.

Args:

  • name (String): The Attribute name.
  • value (String): The Attribute value.

Returns:

An internal representation of an Attribute.

graph

graph(self: Self, name: String, in_types: TypeTuple, out_types: TypeTuple) -> Graph

Adds an empty Graph to this Module.

Args:

  • name (String): The Graph's name.
  • in_types (TypeTuple): The Graphs's input types.
  • out_types (TypeTuple): The Graphs's return types.

Returns:

An empty Graph ready to be filled with ops.

graph(self: Self, in_types: TypeTuple, out_types: TypeTuple) -> Graph

Adds an empty Graph with a default name to this Module.

Args:

  • in_types (TypeTuple): The Graphs's input types.
  • out_types (TypeTuple): The Graphs's return types.

Returns:

An empty Graph ready to be filled with ops.