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

Example

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

Local Variables¶

>>> output
tensor([0.0000, 1.0000, 2.0000], dtype=float32, loc=gpu:0, shape=(3,))