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 trait

TensorStorage

Defines a non-owning interface for accessing tensor storage.

A conforming type describes how to access storage that is owned elsewhere. It provides a concrete StorageType handle along with static operations to load from, store into, and reinterpret values of that handle. The trait never owns the underlying memory; the handle's origin parameter tracks the lifetime and mutability of the borrowed storage.

Implemented traits

AnyType

`comptime` members

`element_size`

comptime element_size = Int(1)

`OffsetResultType`

comptime OffsetResultType[offset_types: TypeList[values]]

The storage type produced by offsetting with a given coordinate.

Parameters

offset_types (TypeList[values]): The coordinate element types of the applied offset.

`StorageType`

comptime StorageType[mut: Bool, //, dtype: DType, origin: Origin[mut=mut], address_space: AddressSpace]

The concrete, register-passable handle to the borrowed storage.

Every operation in this trait acts on values of this type. It is parameterized on the element dtype, the origin that tracks the lifetime and mutability of the borrowed storage, and the address_space the storage resides in, so a single conforming type describes a whole family of handles.

Parameters

mut (Bool): The mutability of the borrowed storage, inferred from origin.
dtype (DType): The element data type of the borrowed storage.
origin (Origin[mut=mut]): The origin tracking the lifetime of the borrowed storage.
address_space (AddressSpace): The address space the borrowed storage resides in.

Required methods

`unsafe_cast`

static def unsafe_cast[to_mut: Bool, //, to_dtype: DType, to_origin: Origin[mut=to_mut], to_address_space: AddressSpace](storage: Self.StorageType[origin.mut, origin_of(origin), dtype, origin, address_space]) -> Self.StorageType[to_mut, origin_of(to_origin), to_dtype, to_origin, to_address_space]

Reinterprets a storage handle with new type parameters.

This performs an unchecked reinterpretation of the underlying reference; no conversion of the stored elements takes place. The caller is responsible for ensuring the new dtype, origin, and address_space are valid for the referenced storage.

Parameters:

to_mut (Bool): The mutability to reinterpret the storage as.
to_dtype (DType): The element data type to reinterpret the storage as.
to_origin (Origin[mut=to_mut]): The origin to reinterpret the storage as.
to_address_space (AddressSpace): The address space to reinterpret the storage as.

Args:

storage (_Self.StorageType[origin.mut, origin_of(origin), dtype, origin, address_space]): The storage to reinterpret.

Returns:

_Self.StorageType[to_mut, origin_of(to_origin), to_dtype, to_origin, to_address_space]: A handle referring to the same storage, viewed with the new type parameters.

`load`

static def load[dtype: DType, //, width: SIMDLength, alignment: Int, invariant: Bool = False, non_temporal: Bool = False](storage: Self.StorageType[False, origin_of(origin), dtype, origin, address_space]) -> SIMD[dtype, width]

Loads a SIMD value from the storage.

Parameters:

dtype (DType): The element data type of the storage.
width (SIMDLength): The number of elements to load.
alignment (Int): The alignment guarantee for the load.
invariant (Bool): If True, the compiler may assume the memory won't be modified during the kernel, enabling load hoisting and caching.
non_temporal (Bool): If True, indicates the data will not be reused soon, allowing the hardware to bypass caches (e.g., streaming loads).

Args:

storage (_Self.StorageType[False, origin_of(origin), dtype, origin, address_space]): The storage to load from.

Returns:

SIMD[dtype, width]: The loaded SIMD value.

Loads a SIMD value at a scalar-element offset from the storage.

Parameters:

dtype (DType): The element data type of the storage.
width (SIMDLength): The number of elements to load.
alignment (Int): The alignment guarantee for the load.
invariant (Bool): If True, the compiler may assume the memory won't be modified during the kernel, enabling load hoisting and caching.
non_temporal (Bool): If True, indicates the data will not be reused soon, allowing the hardware to bypass caches (e.g., streaming loads).

Args:

storage (_Self.StorageType[False, origin_of(origin), dtype, origin, address_space]): The storage to load from.
offset (T): The scalar-element offset to load at.

Returns:

SIMD[dtype, width]: The loaded SIMD value.

`store`

static def store[dtype: DType, alignment: Int, *, non_temporal: Bool = False](storage: Self.StorageType[True, origin_of(origin), dtype, origin, address_space], value: SIMD[dtype])

Stores a SIMD value into the storage.

Parameters:

dtype (DType): The element data type of the storage.
alignment (Int): The alignment guarantee for the store.
non_temporal (Bool): If True, indicates the data will not be reused soon, allowing the hardware to bypass caches (e.g., streaming stores).

Args:

storage (_Self.StorageType[True, origin_of(origin), dtype, origin, address_space]): The storage to store into.
value (SIMD[dtype]): The SIMD value to store.

static def store[dtype: DType, alignment: Int, *, non_temporal: Bool = False](storage: Self.StorageType[True, origin_of(origin), dtype, origin, address_space], offset: T, value: SIMD[dtype])

Stores a SIMD value at a scalar-element offset in the storage.

The caller is responsible for ensuring the storage is actually mutable. The dtype, origin, and address_space are inferred from the storage argument for concrete storage types; callers using the trait through an abstract TensorStorage bound must pass them explicitly (before alignment).

Parameters:

dtype (DType): The element data type of the storage.
alignment (Int): The alignment guarantee for the store.
non_temporal (Bool): If True, indicates the data will not be reused soon, allowing the hardware to bypass caches (e.g., streaming stores).

Args:

storage (_Self.StorageType[True, origin_of(origin), dtype, origin, address_space]): The storage to store into.
offset (T): The scalar-element offset to store at.
value (SIMD[dtype]): The SIMD value to store.

`offset`

static def offset[offset_mut: Bool, offset_types: TypeList[offset_types.values], //, offset_dtype: DType, offset_origin: Origin[mut=offset_mut], offset_address_space: AddressSpace](var storage: Self.StorageType[offset_mut, origin_of(offset_origin), offset_dtype, offset_origin, offset_address_space], var offset_coord: Coord[offset_types]) -> Self.OffsetResultType[offset_types.values, offset_types].StorageType[offset_mut, origin_of(offset_origin), offset_dtype, offset_origin, offset_address_space]

Returns a storage handle offset by a number of scalar elements.

Parameters:

offset_mut (Bool): The mutability of the storage, inferred from offset_origin.
offset_types (TypeList[offset_types.values]): The coordinate element types of offset_coord.
offset_dtype (DType): The element data type of the storage.
offset_origin (Origin[mut=offset_mut]): The origin tracking the lifetime of the storage.
offset_address_space (AddressSpace): The address space the storage resides in.

Args:

storage (_Self.StorageType[offset_mut, origin_of(offset_origin), offset_dtype, offset_origin, offset_address_space]): The storage to offset from.
offset_coord (Coord[offset_types]): A rank-1 coordinate holding the number of scalar elements to advance the handle by.

Returns:

_Self.OffsetResultType[offset_types.values, offset_types].StorageType[offset_mut, origin_of(offset_origin), offset_dtype, offset_origin, offset_address_space]: A handle of the same type starting the given number of scalar elements into the referenced storage.

`copy_from`

static def copy_from[SelfLayoutType: TensorLayout, self_origin: MutOrigin, self_address_space: AddressSpace, OtherLayoutType: TensorLayout, other_mut: Bool, other_origin: Origin[mut=other_mut], other_address_space: AddressSpace, //, dst_dtype: DType, src_dtype: DType, OtherStorage: TensorStorage](storage: Tuple[Self.StorageType[True, origin_of(self_origin), dst_dtype, self_origin, self_address_space], SelfLayoutType], other: Tuple[OtherStorage.StorageType[other_mut, origin_of(other_origin), src_dtype, other_origin, other_address_space], OtherLayoutType])

Copies the elements of other into storage, in place.

Performs an element-by-element copy from other into storage, respecting the layouts of both operands. Each logical element is loaded from other using its layout and stored into storage using its own layout, so the copy works correctly even when the two sides have different shapes or strides (as long as they agree on total element count). When both operands have fully static, row-major layouts and a scalar logical element, the copy widens to SIMD load + cast + SIMD store using the narrower of the two dtypes' native SIMD widths.

Both operands must have statically known shapes with matching total element count.
Both operands must have the same logical element size.
Source and destination dtypes may differ; each logical element is cast to the destination dtype.

Parameters:

SelfLayoutType (TensorLayout): The layout type of the destination storage.
self_origin (MutOrigin): The origin of the destination storage.
self_address_space (AddressSpace): The address space of the destination storage.
OtherLayoutType (TensorLayout): The layout type of the source storage.
other_mut (Bool): The mutability of the source storage.
other_origin (Origin[mut=other_mut]): The origin of the source storage.
other_address_space (AddressSpace): The address space of the source storage.
dst_dtype (DType): The element data type of the destination storage.
src_dtype (DType): The element data type of the source storage.
OtherStorage (TensorStorage): The storage policy of the source. May differ from Self as long as the two policies are copy-compatible (same logical element size).

Args:

storage (Tuple[_Self.StorageType[True, origin_of(self_origin), dst_dtype, self_origin, self_address_space], SelfLayoutType]): A tuple of the destination storage (modified in place) and its layout.
other (Tuple[OtherStorage.StorageType[other_mut, origin_of(other_origin), src_dtype, other_origin, other_address_space], OtherLayoutType]): A tuple of the source storage and its layout.

`distance`

static def distance[dtype: DType, address_space: AddressSpace, //](storage: Self.StorageType[False, origin_of(origin), dtype, origin, address_space], other: Self.StorageType[False, origin_of(origin), dtype, origin, address_space]) -> Int

Returns the scalar-element distance from other to storage.

Parameters:

dtype (DType): The storages' DType.
address_space (AddressSpace): The storages' AddressSpace.

Args:

storage (_Self.StorageType[False, origin_of(origin), dtype, origin, address_space]): The storage to measure the distance to.
other (_Self.StorageType[False, origin_of(origin), dtype, origin, address_space]): The storage to measure the distance from.

Returns:

Int: The number of scalar elements separating the two handles. The value is positive when storage is ahead of other and negative when it precedes other.

Provided methods

`write_type_name_to`

static def write_type_name_to(mut writer: T)

Writes the storage type name representation to the writer.

Args:

writer (T): The Writer to output to.

Implemented traits​

comptime members​

element_size​

OffsetResultType​

Parameters​

StorageType​

Parameters​

Required methods​

unsafe_cast​

load​

store​

offset​

copy_from​

distance​

Provided methods​

write_type_name_to​

Implemented traits

`comptime` members

`element_size`

`OffsetResultType`

Parameters

`StorageType`

Parameters

Required methods

`unsafe_cast`

`load`

`store`

`offset`

`copy_from`

`distance`

Provided methods

`write_type_name_to`