Skip to main content

/

Docs

Nightly

v24.6

Log in

Mojo struct

ManagedTensorSlice

@register_passable(trivial) struct ManagedTensorSlice[type: DType, rank: Int]

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.

Implemented traits​

AnyType, CollectionElement, Copyable, ExplicitlyCopyable, Movable, TensorLike, UnknownDestructibility

Methods​

__init__​

__init__(ptr: UnsafePointer[SIMD[type, 1]], slices: InlineArray[Slice, rank], slicer_spec: RuntimeTensorSpec[type, rank]) -> Self

Initializes a ManagedTensorSlice from a pointer, array of slices and tensor spec.

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

__init__(ptr: UnsafePointer[SIMD[type, 1]], 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.

__init__(ptr: UnsafePointer[SIMD[type, 1]], 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.

__getitem__​

__getitem__(self, indices: IndexList[rank]) -> SIMD[type, 1]

Gets the value at the specified indices.

Args:

  • ​indices (IndexList[rank]): The indices of the value to retrieve.

Returns:

The value at the specified indices.

__getitem__(self, *indices: Int) -> SIMD[type, 1]

Gets the value at the specified indices.

Args:

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

Returns:

The value at the specified indices.

__setitem__​

__setitem__(self, *indices: Int, *, val: SIMD[type, 1])

Stores the value at the specified indices.

Args:

  • ​*indices (Int): The indices of the value to store.
  • ​val (SIMD[type, 1]): The value to store.

__setitem__(self, indices: IndexList[rank], val: SIMD[type, 1])

Stores the value at the specified indices.

Args:

  • ​indices (IndexList[rank]): The indices of the value to store.
  • ​val (SIMD[type, 1]): The value to store.

spec​

spec(self) -> TensorSpec

Gets the TensorSpec of this tensor slice, which provides meta-data about the tensor slice.

Returns:

The static TensorSpec for this tensor slice.

shape​

shape(self) -> IndexList[rank]

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

Returns:

The shape of this tensor slice.

dim_size​

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:

The size of the tensor slice in the given dimension.

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:

The size of the tensor slice in the given dimension.

strides​

strides(self) -> IndexList[rank]

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

Returns:

The strides of this tensor slice.

size​

size(self) -> Int

Computes the tensor slice's number of elements.

Returns:

The total number of elements in the tensor slice.

unsafe_ptr​

unsafe_ptr[__type: DType = type](self) -> UnsafePointer[SIMD[__type, 1]]

Get the pointer stored in this tensor slice.

Danger: This method obtains the pointer stored in this tensor slice. In general, it should not be used, as it can modify the invariants of this tensor slice and lead to unexpected behavior. Custom operations should avoid using this method.

Parameters:

  • ​__type (DType): The type of the UnsafePointer in this tensor slice.

Returns:

The UnsafePointer which contains the data for this tensor slice.

load​

load[width: Int, _rank: Int](self, index: IndexList[_rank]) -> SIMD[type, width]

Gets data from this tensor slice as a SIMD.

Danger: This method separates the data of this tensor slice from the tensor slice itself. Custom operations should avoid using this method.

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.

Args:

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

Returns:

Data from this tensor slice at dimension index.

store​

store[width: Int, _rank: Int](self, index: IndexList[_rank], val: SIMD[type, width])

Sets data in this tensor slice from a SIMD.

Danger: This method changes the data in this tensor slice without any safety guarantees. Custom operations should avoid using this method.

Parameters:

  • ​width (Int): The width of the SIMD value.
  • ​_rank (Int): The rank of the tensor slice.

Args:

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