IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

ManagedTensorSlice

struct ManagedTensorSlice[mut: Bool, input: IO, dtype: DType, rank: Int, InFusion: InputFusion, OutFusion: OutputFusion, ComputeFusion: ComputeOutputFusion, ComputeFusionTile: ComputeOutputFusionTile, OutFusionTile: OutputFusionTile, //, io_spec: IOSpec[mut, input], *, static_spec: StaticTensorSpec[dtype, rank, static_spec.static_layout, InFusion, OutFusion, ComputeFusion, ComputeFusionTile, OutFusionTile]]

A view of a tensor that does not own the underlying allocated pointer. When the object lifetime ends it does not free the underlying pointer. Conversely, if a ManagedTensorSlice is created, it will not extend the life of the underlying pointer.

Therefore, the user must take care to keep the pointer alive until the last use of a ManagedTensorSlice instance. This class is useful for writing custom operations where memory is managed by an external runtime like in MAX's inference stack.

Fields​

  • ​in_fusion (InFusion):
  • ​out_fusion (OutFusion):
  • ​compute_fusion (ComputeFusion):
  • ​compute_fusion_tile (ComputeFusionTile):
  • ​output_fusion_tile (OutFusionTile):

Implemented traits​

AnyType, Copyable, Deinitable, DevicePassable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable, Writable

comptime members​

address_space​

comptime address_space = static_spec.address_space

alignment​

comptime alignment = static_spec.alignment

device_type​

comptime device_type = LayoutTensor[dtype, static_spec.to_layout(), MutAnyOrigin]

RuntimeLayout​

comptime RuntimeLayout = Layout[static_layout._shape_types, static_layout._stride_types]

Methods​

__init__​

def __init__(ptr: Pointer[Scalar[dtype]], shape: IndexList[rank]) -> Self

Initializes a ManagedTensorSlice from a pointer and shape.

In general, custom operations should not create ManagedTensorSlice instances, but instead use the ones provided by the MAX inference engine.

def __init__(ptr: Pointer[Scalar[dtype]], shape: IndexList[rank], strides: IndexList[rank]) -> Self

Initializes a ManagedTensorSlice from a pointer, shape, and strides.

In general, custom operations should not create ManagedTensorSlice instances, but instead use the ones provided by the MAX inference engine.

def __init__(ptr: Pointer[Scalar[dtype]], shape: Coord[static_spec.static_layout._shape_types], strides: Coord[static_spec.static_layout._stride_types]) -> Self

Initializes a ManagedTensorSlice from a pointer, shape, and strides.

The shape and strides are provided as Coords; their runtime values fill the dynamic dimensions of the tensor slice's layout.

In general, custom operations should not create ManagedTensorSlice instances, but instead use the ones provided by the MAX inference engine.

__getitem__​

def __getitem__(self, indices: IndexList[rank]) -> Scalar[dtype]

Gets the value at the specified indices.

Args:

Returns:

Scalar[dtype]: The value at the specified indices.

def __getitem__(self, *indices: Int) -> Scalar[dtype]

Gets the value at the specified indices.

Args:

  • ​*indices (Int): The indices of the value to retrieve.

Returns:

Scalar[dtype]: The value at the specified indices.

__setitem__​

def __setitem__(self, *indices: Int, *, val: Scalar[dtype])

Stores the value at the specified indices.

Args:

  • ​*indices (Int): The indices of the value to store.
  • ​val (Scalar[dtype]): The value to store.

def __setitem__(self, indices: IndexList[rank], val: Scalar[dtype])

Stores the value at the specified indices.

Args:

get_type_name​

static def get_type_name() -> String

Returns:

String

shape​

def shape(self) -> IndexList[rank]

Gets the shape of this tensor slice, as an IndexList.

Returns:

IndexList[rank]: The shape of this tensor slice.

shape_coord​

def shape_coord(self) -> Coord[static_spec.static_layout._shape_types]

Gets the shape of this tensor slice as a Coord.

Unlike shape, which returns a runtime IndexList, the returned Coord preserves the static-vs-dynamic structure of the tensor's static layout: statically-known dimensions are encoded as compile-time values in the Coord's type, while dynamic dimensions are filled from the runtime shape.

Returns:

Coord[static_spec.static_layout._shape_types]: The shape of this tensor slice as a Coord.

strides_coord​

def strides_coord(self) -> Coord[static_spec.static_layout._stride_types]

Gets the strides of this tensor slice as a Coord.

Unlike strides, which returns a runtime IndexList, the returned Coord preserves the static-vs-dynamic structure of the tensor's static layout: statically-known strides are encoded as compile-time values in the Coord's type, while dynamic strides are filled from the runtime strides.

Returns:

Coord[static_spec.static_layout._stride_types]: The strides of this tensor slice as a Coord.

runtime_layout​

def runtime_layout(self) -> Self.RuntimeLayout

Gets the runtime layout of this tensor slice.

The layout bundles the shape and strides as Coords, preserving the static-vs-dynamic structure of the tensor's static layout.

Returns:

Self.RuntimeLayout: The runtime layout of this tensor slice.

dim_size​

def dim_size(self, index: Int) -> Int

Gets the size of a given dimension of this tensor slice using a run time value.

Args:

  • ​index (Int): The zero-based index of the dimension.

Returns:

Int: The size of the tensor slice in the given dimension.

def dim_size[index: Int](self) -> Int

Gets the size of a given dimension of this tensor slice using a compile time value.

Parameters:

  • ​index (Int): The zero-based index of the dimension.

Returns:

Int: The size of the tensor slice in the given dimension.

strides​

def strides(self) -> IndexList[rank]

Gets the strides of this tensor slice, as an IndexList.

Returns:

IndexList[rank]: The strides of this tensor slice.

stride_length​

def stride_length(self, index: Int) -> Int

Gets the length of the stride of a given dimension of this tensor slice using a run time value.

Args:

  • ​index (Int): The zero-based index of the dimension.

Returns:

Int: The size of the tensor slice in the given dimension.

def stride_length[index: Int](self) -> Int

Gets the length of the stride of a given dimension of this tensor slice using a compile time value.

Parameters:

  • ​index (Int): The zero-based index of the dimension.

Returns:

Int: The size of the tensor slice in the given dimension.

size​

def size(self) -> Int

Computes the tensor slice's number of elements.

Returns:

Int: The total number of elements in the tensor slice.

bytecount​

def bytecount(self) -> Int

Returns the size of the tensor slice in bytes.

Returns:

Int: The total number of bytes in the tensor slice.

unsafe_ptr​

def unsafe_ptr[_dtype: DType = dtype](self) -> Pointer[Scalar[_dtype], MutAnyOrigin]

Get the pointer stored in this tensor slice.

Since this method obtains the pointer stored in this tensor slice, it can modify the invariants of this tensor slice and lead to unexpected behavior. It should be used with caution.

Parameters:

  • ​_dtype (DType): The type of the Pointer in this tensor slice.

Returns:

Pointer[Scalar[_dtype], MutAnyOrigin]: The Pointer which contains the data for this tensor slice.

to_device_buffer​

def to_device_buffer(self, ctx: DeviceContext) -> DeviceBuffer[dtype]

Returns:

DeviceBuffer[dtype]

load​

def load[width: Int, _rank: Int, element_alignment: Int = Int(1)](self, index: IndexList[_rank]) -> SIMD[dtype, width]

Gets data from this tensor slice as a SIMD.

Parameters:

  • ​width (Int): The width of the SIMD value. This must be large enough to contain the data from this tensor slice.
  • ​_rank (Int): The rank of the tensor slice.
  • ​element_alignment (Int): Indicate the alignment of the pointer stored to memory. This is needed to issue vector load for GPUs with strict alignment requirements.

Args:

  • ​index (IndexList[_rank]): An IndexList of size _rank to indicate the dimension of the tensor slice to obtain data from.

Returns:

SIMD[dtype, width]: Data from this tensor slice at dimension index.

def load[width: Int, element_alignment: Int = Int(1)](self, index: Coord) -> SIMD[dtype, width]

Gets data from this tensor slice as a SIMD, indexed by a Coord.

Parameters:

  • ​width (Int): The width of the SIMD value. This must be large enough to contain the data from this tensor slice.
  • ​element_alignment (Int): Indicate the alignment of the pointer stored to memory. This is needed to issue vector load for GPUs with strict alignment requirements.

Args:

  • ​index (Coord): A Coord indicating the dimension of the tensor slice to obtain data from.

Returns:

SIMD[dtype, width]: Data from this tensor slice at dimension index.

store​

def store[width: SIMDLength, _rank: Int, element_alignment: Int = Int(1)](self: ManagedTensorSlice[static_spec=static_spec], index: IndexList[_rank], val: SIMD[dtype, width])

Sets data in this tensor slice from a SIMD.

Parameters:

  • ​width (SIMDLength): The width of the SIMD value.
  • ​_rank (Int): The rank of the tensor slice.
  • ​element_alignment (Int): Indicate the alignment of the pointer stored to memory. This is needed to issue vector store for GPUs with strict alignment requirements.

Args:

  • ​index (IndexList[_rank]): An IndexList of size _rank to indicate the dimension of the tensor slice to set data in.
  • ​val (SIMD[dtype, width]): The data to set into this tensor slice.

def store[width: SIMDLength, element_alignment: Int = Int(1)](self: ManagedTensorSlice[static_spec=static_spec], index: Coord, val: SIMD[dtype, width])

Sets data in this tensor slice from a SIMD, indexed by a Coord.

Parameters:

  • ​width (SIMDLength): The width of the SIMD value.
  • ​element_alignment (Int): Indicate the alignment of the pointer stored to memory. This is needed to issue vector store for GPUs with strict alignment requirements.

Args:

  • ​index (Coord): A Coord indicating the dimension of the tensor slice to set data in.
  • ​val (SIMD[dtype, width]): The data to set into this tensor slice.

with_tile_layout​

def with_tile_layout[new_layout: TensorLayout](self, new_runtime_shape: IndexList[new_layout.rank], new_runtime_strides: IndexList[new_layout.rank], offset_ptr: Optional[Pointer[Scalar[dtype], MutAnyOrigin]] = None) -> ManagedTensorSlice[io_spec, static_spec=static_spec.with_tile_layout[new_layout]()]

Returns:

ManagedTensorSlice[io_spec, static_spec=static_spec.with_tile_layout[new_layout]()]

to_layout_tensor​

def to_layout_tensor(self) -> LayoutTensor[dtype, static_spec.to_layout(), MutAnyOrigin]

Returns:

LayoutTensor[dtype, static_spec.to_layout(), MutAnyOrigin]

to_tile_tensor​

def to_tile_tensor[coord_dtype: DType = DType.int64](self) -> TileTensor[dtype, Layout[static_spec.static_layout._shape_types, static_spec.static_layout._stride_types], MutUntrackedOrigin]

Returns:

TileTensor[dtype, Layout[static_spec.static_layout._shape_types, static_spec.static_layout._stride_types], MutUntrackedOrigin]

write_to​

def write_to(self, mut writer: T)

Formats this buffer to the provided Writer.

Args:

  • ​writer (T): The object to write to.

write_repr_to​

def write_repr_to(self, mut writer: T)

Formats this buffer to the provided Writer.

Args:

  • ​writer (T): The object to write to.

Models

Cookbook

Puzzles

Blog

Community

Contact