iota

nvtripy.iota(shape: Sequence[int | DimensionSize], dim: int = 0, dtype: dtype = float32) → Tensor[source]

Fills an output tensor with consecutive values starting from zero along the given dimension.

Parameters:

  • shape (Sequence[int | DimensionSize]) – The desired shape.

  • dim (int) – Dimension along which to perform the iota operation. This cannot exceed the rank of the specified shape.

  • dtype (dtype) – The desired data type.

Returns:

A tensor of shape shape and data type dtype.

Return type:

Tensor

INPUT REQUIREMENTS:

dtype is one of [float32, float16, bfloat16, float8, int4, int8, int32, int64, bool]

OUTPUT GUARANTEES:

return[0].dtype == dtype

Example
1output = tp.iota((3,), dim=-1)
Local Variables
>>> output
tensor([0, 1, 2], dtype=float32, loc=gpu:0, shape=(3,))