Skip to main content

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?