iota_like¶

tripy.iota_like(input: Tensor, dim: int = 0, dtype: dtype | None = None) → Tensor[source]¶

Returns a tensor of the same shape and data type as the input tensor, with consecutive values starting from zero along the given dimension.

Parameters:

input (Tensor) – Input tensor.
dim (int) – Dimension along which to perform the iota operation. This cannot exceed the rank of the specified shape.
dtype (dtype | None) – The desired data type. This will override the data type inferred from the input tensor.

Returns:

A tensor of the same shape and data type (unless dtype is provided) as the input.

Return type:

Tensor

Example

Example¶

input = tp.Tensor([1, 2, 3])
output = tp.iota_like(input)

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