PyTorch `iota` operator

iota.py

from torch import Tensor
import torch


def iota(shape: tuple[int, ...], dim: int = 0, **kwargs) -> Tensor:
    dimensions = []

    for index, _ in enumerate(shape):
        if index != dim:
            dimension = 1
        else:
            dimension = shape[index]

        dimensions = [*dimensions, dimension]

    return torch.arange(shape[dim], **kwargs).view(*dimensions).expand(*shape)

iota([4, 8], 0, dtype=torch.int32)

tensor([[0, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1],
        [2, 2, 2, 2, 2, 2, 2, 2],
        [3, 3, 3, 3, 3, 3, 3, 3]], dtype=torch.int32)

iota([4, 8], 1, dtype=torch.int32)

tensor([[0, 1, 2, 3, 4, 5, 6, 7],
        [0, 1, 2, 3, 4, 5, 6, 7],
        [0, 1, 2, 3, 4, 5, 6, 7],
        [0, 1, 2, 3, 4, 5, 6, 7]], dtype=torch.int32)

iota([2, 4, 8], 0, dtype=torch.int32)

tensor([[[0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0, 0, 0, 0]],

        [[1, 1, 1, 1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1, 1, 1, 1],
         [1, 1, 1, 1, 1, 1, 1, 1]]], dtype=torch.int32)