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,
ImplicitlyDestructible,
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
Returns the compile-time value of the i-th shape dimension.
static_stride
comptime static_stride
Returns the compile-time value of the i-th stride dimension.
stride_known
comptime stride_known
Whether all stride dimensions are known at compile time.
Required methods
__init__
__init__(out self: _Self, *, copy: _Self)
Create a new instance of the value by copying an existing one.
Args:
- copy (
_Self): The value to copy.
Returns:
_Self
__init__(out self: _Self, *, deinit take: _Self)
Create a new instance of the value by moving the value of another.
Args:
- take (
_Self): The value to move.
Returns:
_Self
shape
shape[i: Int](self: _Self) -> _Self._shape_types[i]
Returns the i-th shape dimension.
Parameters:
- i (
Int): The dimension index.
Returns:
_Self._shape_types: The shape value for dimension i.
stride
stride[i: Int](self: _Self) -> _Self._stride_types[i]
Returns the i-th stride dimension.
Parameters:
- i (
Int): The dimension index.
Returns:
_Self._stride_types: The stride value for dimension i.
product
product(self: _Self) -> Int
Returns the total number of elements in the layout's domain.
Returns:
Int: The product of all shape dimensions.
size
size(self: _Self) -> Int
Returns the total number of elements. Alias for product().
Returns:
Int: The product of all shape dimensions.
__call__
__call__[index_type: CoordLike, *, linear_idx_type: DType = DType.int64](self: _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: The linear memory index corresponding to the given coordinates.
idx2crd
idx2crd[*, out_dtype: DType = DType.int64](self: _Self, idx: Int) -> Coord[#kgen.variadic.tabulate(_Self.rank, [idx: __mlir_type.index] RuntimeInt[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.
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: A Coord containing the logical coordinates corresponding to
the linear index.
shape_coord
shape_coord(self: _Self) -> Coord[_Self._shape_types]
Returns the full shape as a Coord.
Returns:
Coord: A Coord containing all shape dimensions.
stride_coord
stride_coord(self: _Self) -> Coord[_Self._stride_types]
Returns the full stride as a Coord.
Returns:
Coord: A Coord containing all stride dimensions.
transpose
transpose(self: _Self) -> Layout[#kgen.variadic.reduce(_Self._shape_types, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, VA[(add (mul idx, -1), len(VA), -1)])), #kgen.variadic.reduce(_Self._stride_types, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, VA[(add (mul idx, -1), len(VA), -1)]))]
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
make_dynamic[dtype: DType](self: _Self) -> Layout[#kgen.variadic.reduce(_Self._shape_types, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, RuntimeInt[dtype])), #kgen.variadic.reduce(_Self._stride_types, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, RuntimeInt[dtype]))]
Converts all dimensions to runtime values of the given dtype.
Parameters:
- dtype (
DType): The data type for the resultingRuntimeIntvalues.
Returns:
Layout: A new Layout with all dimensions as RuntimeInt[dtype].
Provided methods
copy
copy(self: _Self) -> _Self
Explicitly construct a copy of self, a convenience method for Self(copy=self) when the type is inconvenient to write out.
Returns:
_Self: A copy of this value.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!