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

TensorLayout

Trait defining the interface for mixed compile-time/runtime layouts.

Implementors map logical multi-dimensional coordinates to linear memory indices, with support for dimensions that are known at compile time or determined at runtime.

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

all_dims_known

comptime all_dims_known = _Self.shape_known and _Self.stride_known

Whether all shape and stride dimensions are known at compile time.

flat_rank

comptime flat_rank

The number of dimensions after flattening nested coordinates.

rank

comptime rank

The number of dimensions in the layout.

shape_known

comptime shape_known

Whether all shape dimensions are known at compile time.

static_cosize

comptime static_cosize

The compile-time size of the memory region spanned by the layout.

static_product

comptime static_product

The compile-time product of all shape dimensions.

static_shape

comptime static_shape[i: Int]

Returns the compile-time value of the i-th shape dimension.

Parameters

• ​i (Int): The dimension index.

static_stride

comptime static_stride[i: Int]

Returns the compile-time value of the i-th stride dimension.

Parameters

• ​i (Int): The dimension index.

stride_known

comptime stride_known

Whether all stride dimensions are known at compile time.

Required methods

__init__

def __init__(out self, *, copy: Self)

Create a new instance of the value by copying an existing one.

Args:

• ​copy (_Self): The value to copy.

Returns:

_Self

def __init__(out self, *, deinit move: Self)

Create a new instance of the value by moving the value of another.

Args:

• ​move (_Self): The value to move.

Returns:

_Self

shape

def shape[i: Int](self) -> _Self.__shape_types[SIMDLength(i)]

Returns the i-th shape dimension.

Parameters:

• ​i (Int): The dimension index.

Returns:

_Self.__shape_types[SIMDLength(i)]: The shape value for dimension i.

stride

def stride[i: Int](self) -> _Self.__stride_types[SIMDLength(i)]

Returns the i-th stride dimension.

Parameters:

• ​i (Int): The dimension index.

Returns:

_Self.__stride_types[SIMDLength(i)]: The stride value for dimension i.

product

def product(self) -> Int

Returns the total number of elements in the layout's domain.

Returns:

Int: The product of all shape dimensions.

size

def size(self) -> Int

Returns the total number of elements. Alias for product().

Returns:

Int: The product of all shape dimensions.

__call__

def __call__[index_type: CoordLike, *, linear_idx_type: DType = DType.int64](self, index: index_type) -> Scalar[linear_idx_type]

Maps a logical coordinate to a linear memory index.

Parameters:

• ​index_type (CoordLike): The coordinate type.
• ​linear_idx_type (DType): The data type for the returned linear index.

Args:

• ​index (index_type): The logical coordinates to map.

Returns:

Scalar[linear_idx_type]: The linear memory index corresponding to the given coordinates.

idx2crd

def idx2crd[*, out_dtype: DType = DType.int64](self, idx: Int) -> Coord[*#kgen.param_list.tabulate(len(_Self.__shape_types), [idx: __mlir_type.index] Coord[*#kgen.param_list.tabulate(len(_Self.__shape_types[SIMDLength(Int(idx))]._ParamListType), [idx: __mlir_type.index] Scalar[out_dtype])] if _Self.__shape_types[SIMDLength(Int(idx))].is_tuple else Scalar[out_dtype])]

Maps a linear memory index back to logical coordinates.

This is the inverse of __call__ (crd2idx). Given a linear index, it computes the corresponding multi-dimensional coordinates. For hierarchical layouts (e.g. from zipped_divide), the result preserves the nested coordinate structure.

Examples: For a layout with shape (3, 4) and row-major strides:

• layout.idx2crd(0) returns (0, 0).
• layout.idx2crd(5) returns (1, 1).
• layout.idx2crd(11) returns (2, 3).

Parameters:

• ​out_dtype (DType): The data type for the output coordinate values.

Args:

• ​idx (Int): The linear memory index to convert to coordinates.

Returns:

Coord[*#kgen.param_list.tabulate(len(_Self.__shape_types), [idx: __mlir_type.index] Coord[*#kgen.param_list.tabulate(len(_Self.__shape_types[SIMDLength(Int(idx))]._ParamListType), [idx: __mlir_type.index] Scalar[out_dtype])] if _Self.__shape_types[SIMDLength(Int(idx))].is_tuple else Scalar[out_dtype])]: A Coord containing the logical coordinates corresponding to the linear index. For nested layouts, the result mirrors the shape nesting with Scalar leaves.

shape_coord

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

Returns the full shape as a Coord.

Returns:

Coord[_Self._shape_types]: A Coord containing all shape dimensions.

stride_coord

def stride_coord(self) -> Coord[_Self._stride_types]

Returns the full stride as a Coord.

Returns:

Coord[_Self._stride_types]: A Coord containing all stride dimensions.

transpose

def transpose(self) -> Layout[*?, *?]

Transposes the layout by reversing the order of dimensions.

For an n-dimensional layout, this reverses the order of both shapes and strides. For 2D layouts, this swaps rows and columns.

Returns:

Layout[*?, *?]: A new Layout with transposed dimensions.

make_dynamic

def make_dynamic[dtype: DType](self) -> Layout[*?, *?]

Converts all dimensions to runtime values of the given dtype.

Parameters:

• ​dtype (DType): The data type for the resulting Scalar values.

Returns:

Layout[*?, *?]: A new Layout with all dimensions as Scalar[dtype].

Provided methods

copy

def copy(self) -> Self

Explicitly construct a copy of self, a convenience method for Self(copy=self) when the type is inconvenient to write out.

Overriding this method is not allowed.

Returns:

_Self: A copy of this value.