Mojo struct
ManagedTensorSlice
@register_passable(trivial)
struct ManagedTensorSlice[type: DType, rank: Int]
ManagedTensorSlice is like TensorSlice but it does not effect the life of the underlying allocated pointer. Unlike TensorSlice, when the object lifetime ends it does not affect the lifetime of the underlying pointer. Conversly, if a ManagedTensorSlice is created, it will not extend the life of the underlying pointer.
Therefore, the user must take care to keep the ptr alive until ManagedTensorSlice's last use. This class is useful for writing kernels 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
__init__(ptr: UnsafePointer[SIMD[type, 1]], shape: IndexList[rank]) -> Self
__init__(ptr: UnsafePointer[SIMD[type, 1]], shape: IndexList[rank], strides: IndexList[rank]) -> Self
@implicit
__init__(ndbuffer: NDBuffer[type, rank]) -> Self
Initializes a ManagedTensorSlice from an NDBuffer.
Note that forwarding of static shape, strides, and lambdas won't work.
__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
shape
shape(self) -> IndexList[rank]
dim_size
dim_size(self, index: Int) -> Int
dim_size[index: Int](self) -> Int
strides
strides(self) -> IndexList[rank]
stride_length
stride_length(self, index: Int) -> Int
stride_length[index: Int](self) -> Int
size
size(self) -> Int
Computes the tensor slice's number of elements.
Returns:
The total number of elements in the ManagedTensorSlice.
unsafe_ptr
unsafe_ptr[__type: DType = type](self) -> UnsafePointer[SIMD[__type, 1]]
load
load[width: Int, _rank: Int](self, index: IndexList[_rank]) -> SIMD[type, width]
store
store[width: Int, _rank: Int](self, index: IndexList[_rank], val: SIMD[type, width])
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!