Mojo function
idx2crd
idx2crd[Shape: CoordLike, Stride: CoordLike, out_dtype: DType = DType.int64](idx: Int, shape: Shape, stride: Stride) -> Coord[#kgen.variadic.reduce(Shape.VariadicType, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, ComptimeInt[0] if (VA[idx].static_value == 1) if VA[idx].is_static_value else VA[idx].is_static_value else RuntimeInt[out_dtype]))]
Calculate the coordinate tuple from a linear index.
This is the inverse of crd2idx - given a linear index, shape, and stride,
it computes the multi-dimensional coordinates using the per-element formula:
coord[i] = (idx // stride[i]) % shape[i].
When a shape dimension is statically known to be 1, the corresponding output coordinate is a ComptimeInt[0]. Otherwise, coordinates are RuntimeInt[out_dtype].
The idx, and all components of shape and stride must be non-negative.
Examples: For a 2D tensor with shape (3, 4) and row-major strides (4, 1):
- idx2crd(0, shape, stride) returns (0, 0).
- idx2crd(5, shape, stride) returns (1, 1).
- idx2crd(11, shape, stride) returns (2, 3).
Parameters:
- Shape (
CoordLike): The shape type (must be CoordLike). - Stride (
CoordLike): The stride type (must be CoordLike). - out_dtype (
DType): The output data type for coordinate values.
Args:
- idx (
Int): The linear index to convert. - shape (
Shape): The shape of the tensor. - stride (
Stride): The stride of the tensor.
Returns:
Coord: A Coord containing the coordinate values for each dimension.
Dimensions with static shape 1 produce ComptimeInt[0], others
produce RuntimeInt[out_dtype].
idx2crd[Index: CoordLike, Shape: CoordLike, Stride: CoordLike, out_dtype: DType = DType.int64](idx: Index, shape: Shape, stride: Stride) -> Coord[#kgen.variadic.reduce(Shape.VariadicType, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, ComptimeInt[0] if (VA[idx].static_value == 1) if VA[idx].is_static_value else VA[idx].is_static_value else ComptimeInt[((Index.static_value // Stride.VariadicType[idx].static_value) % VA[idx].static_value)] if Stride.VariadicType[idx].is_static_value if VA[idx].is_static_value if Index.is_static_value else Index.is_static_value else VA[idx].is_static_value if Index.is_static_value else Index.is_static_value else RuntimeInt[out_dtype]))]
Calculate the coordinate tuple from a CoordLike linear index.
This overload accepts a CoordLike index, enabling compile-time result
computation when the index, shape, and stride are all statically known.
Uses the per-element formula: coord[i] = (idx // stride[i]) % shape[i].
The idx, and all components of shape and stride must be non-negative.
Parameters:
- Index (
CoordLike): The index type (must be CoordLike). - Shape (
CoordLike): The shape type (must be CoordLike). - Stride (
CoordLike): The stride type (must be CoordLike). - out_dtype (
DType): The output data type for coordinate values.
Args:
- idx (
Index): The CoordLike linear index to convert. - shape (
Shape): The shape of the tensor. - stride (
Stride): The stride of the tensor.
Returns:
Coord: A Coord containing the coordinate values for each dimension.
When idx, shape, and stride are all compile-time known, produces
ComptimeInt results. Otherwise produces RuntimeInt[out_dtype].
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!