iota

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

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) – [dtype=T1] The desired data type.

Returns:

[dtype=T1] A tensor of shape shape and data type dtype.

Return type:

Tensor

DATA TYPE CONSTRAINTS:

• T1: bool, float8, float16, bfloat16, float32, int4, int8, int32, int64

Example

output = tp.iota((3,), dim=-1)

Local Variables

>>> output
tensor([0, 1, 2], dtype=float32, loc=gpu:0, shape=(3,))