Skip to main content

session

The entry point to MAX Engine and can be used to load models for inference.

LoadOptions

Configuration options to load model with MAX Engine.

If you're loading a TorchScript model, you must create an instance of this type and then call add_input_spec() or add_input_specs(), depending on whether your model has one or multiple inputs (you must specify an input spec for every input). Then, pass this LoadOptions along with your TorchScript model to Session.load_model().

Implemented traits:

AnyType, CollectionElement, Copyable, Movable

Methods:

__init__

__init__(inout self: Self)

Creates a new LoadOptions object.

set_custom_ops_path

set_custom_ops_path(inout self: Self, path: Path)

Replace Modular kernels in given model with user-defined kernels.

Args:

  • path (Path): Path to mojo custom op package.

add_input_spec

add_input_spec(inout self: Self, spec: TensorSpec)

Add specifications for one input tensor, as a TensorSpec. Only applicable for TorchScript models.

If an input supports dynamic shapes, use None for that dimension size in the TensorSpec.

Args:

  • spec (TensorSpec): Spec for the input. This is the standard library TensorSpec.

add_input_spec(inout self: Self, shape: Optional[List[Optional[SIMD[si64, 1]]]], dtype: DType)

Add specifications for one input tensor, as a list of integers. Only applicable for TorchScript models.

If an input supports dynamic shapes, use None for that dimension size.

Args:

  • shape (Optional[List[Optional[SIMD[si64, 1]]]]): Shape of the input, as a list of integers.
  • dtype (DType): Datatype of the input, from the standard library DType.

add_input_specs

add_input_specs(inout self: Self, specs: List[TensorSpec])

Add specifications for multiple input tensors, as a list of TensorSpec objects. Only applicable for TorchScript models.

If an input supports dynamic shapes, use None for that dimension size in the TensorSpec.

Args:

  • specs (List[TensorSpec]): A list of specs for the inputs. This is the standard library TensorSpec.

add_input_specs(inout self: Self, shapes: List[Optional[List[Optional[SIMD[si64, 1]]]]], dtypes: InlinedFixedVector[DType, _calculate_fixed_vector_default_size[AnyRegType][DType]()])

Add specifications for multiple input tensors, using lists of integers. Only applicable for TorchScript models.

If an input supports dynamic shapes, use None for that dimension size.

Args:

  • shapes (List[Optional[List[Optional[SIMD[si64, 1]]]]]): A list of input shapes, in which each input shape is specified as a list if integers.
  • dtypes (InlinedFixedVector[DType, _calculate_fixed_vector_default_size[AnyRegType][DType]()]): Datatypes of the inputs, as a collection of DType values in an InlinedFixedVector.

SessionOptions

Configuration options for InferenceSession.

Implemented traits:

AnyType, Copyable, Movable

Methods:

__init__

__init__(inout self: Self)

Creates a new SessionOptions object.

InferenceSession

Holds the context of MAX Engine and can be used for loading models.

Implemented traits:

AnyType, Copyable, Movable

Methods:

__init__

__init__(inout self: Self, options: SessionOptions)

Creates a new inference session.

Args:

  • options (SessionOptions): Session options to configure how session is created.

load_model

load_model(self: Self, path: Path, config: Optional[LoadOptions]) -> Model

Compile and initialize a model in MAX Engine, with the given model path and config.

Note: PyTorch models must be in TorchScript format, and TensorFlow models must be in SavedModel format. Or pass any ONNX model.

If you're loading a TorchScript model, you must specify the config argument with an instance of LoadOptions that specifies the model's input specs (which may have dynamic shapes).

Args:

  • path (Path): Location of model in filesystem.
  • config (Optional[LoadOptions]): Configurations need for compiling model.

Returns:

Initialized model ready for inference.

load_model(self: Self, graph: Graph, config: Optional[LoadOptions]) -> Model

Compile and initialize a model in MAX Engine, with the given Graph and config.

Args:

  • graph (Graph): MAX Graph module.
  • config (Optional[LoadOptions]): Configurations need for compiling model.

Returns:

Initialized model ready for inference.

load_model(self: Self, module: Module, config: Optional[LoadOptions]) -> Model

Compile and initialize a model in MAX Engine, with the given Module and config.

Args:

  • module (Module): MAX Graph module.
  • config (Optional[LoadOptions]): Configurations need for compiling model.

Returns:

Initialized model ready for inference.

get_as_engine_tensor_spec

get_as_engine_tensor_spec(self: Self, name: String, spec: TensorSpec) -> EngineTensorSpec

Gets a TensorSpec compatible with MAX Engine.

Args:

  • name (String): Name of the Tensor.
  • spec (TensorSpec): Tensor specification in Mojo TensorSpec format.

Returns:

EngineTensorSpec to be used with MAX Engine APIs.

get_as_engine_tensor_spec(self: Self, name: String, shape: Optional[List[Optional[SIMD[si64, 1]]]], dtype: DType) -> EngineTensorSpec

Gets a TensorSpec compatible with MAX Engine.

Args:

  • name (String): Name of the Tensor.
  • shape (Optional[List[Optional[SIMD[si64, 1]]]]): Shape of the Tensor. Dynamic Dimensions can be represented with None and for Dynamic Rank Tensor use None as value for shape.
  • dtype (DType): DataType of the Tensor.

Returns:

EngineTensorSpec to be used with MAX Engine APIs.

new_tensor_map

new_tensor_map(self: Self) -> TensorMap

Gets a new TensorMap. This can be used to pass inputs to model.

Returns:

A new instance of TensorMap.

new_borrowed_tensor_value

new_borrowed_tensor_value[type: DType](self: Self, tensor: Tensor[type]) -> Value

Create a new Value representing data borrowed from given tensor.

The user must ensure the tensor stays live through the lifetime of the value.

Parameters:

  • type (DType): Data type of the tensor to turn into a Value.

Args:

  • tensor (Tensor[type]): Tensor to borrow into a value.

Returns:

A value borrowing the tensor.

new_bool_value

new_bool_value(self: Self, value: Bool) -> Value

Create a new Value representing a Bool.

Args:

  • value (Bool): Boolean to wrap into a value.

Returns:

Value representing the given boolean.

new_list_value

new_list_value(self: Self) -> Value

Create a new Value representing an empty list.

Returns:

A new value containing an empty list.