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
__copy_ctor_is_trivial
comptime __copy_ctor_is_trivial
A flag (often compiler generated) to indicate whether the implementation of the copy constructor is trivial.
A copy constructor is considered to be trivial if:
- The struct has a compiler-generated trivial copy constructor because all its fields have trivial copy constructors.
In practice, it means the value can be copied by copying the bits from one location to another without side effects.
__del__is_trivial
comptime __del__is_trivial
A flag (often compiler generated) to indicate whether the implementation of __del__ is trivial.
The implementation of __del__ is considered to be trivial if:
- The struct has a compiler-generated trivial destructor and all its fields
have a trivial
__del__method.
In practice, it means that the __del__ can be considered as no-op.
__move_ctor_is_trivial
comptime __move_ctor_is_trivial
A flag (often compiler generated) to indicate whether the implementation of move constructor is trivial.
The implementation of a move constructor is considered to be trivial if:
- The struct has a compiler-generated trivial move constructor because all its fields have trivial move constructors.
In practice, it means the value can be moved by moving the bits from one location to another without side effects.
all_dims_known
comptime all_dims_known = _Self.stride_known if _Self.shape_known else _Self.shape_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_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._mlir_value]
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._mlir_value]
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.splat(RuntimeInt[out_dtype], _Self.rank._mlir_value)]
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.
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!