mi.TensorXf.unsqueeze()

[maxfrei750]

It would be nice to be able to unsqueeze Tensors like in PyTorch.

I created some code and would like to know

1. if there are better ways
2. if it would make sense to add this functionality to Mitsuba

    def unsqueeze(a: mi.TensorXf, dimension: int = 0) -> mi.TensorXf:
        """Add a dimension of size one to a tensor.

        Args:
            a (mi.TensorXf): Tensor to add a dimension to.
            dimension (int, optional): Dimension to add. Defaults to 0.

        Raises:
            ValueError: If the dimension is out of range [-a.ndim - 1, a.ndim + 1).

        Returns:
            mi.TensorXf: Tensor with an additional dimension.
        """

        num_dimensions_original: int = a.ndim    # type: ignore[assignment]

        # Raise an error, if the dimension is not in [-a.ndim - 1, a.ndim + 1).
        if not (-num_dimensions_original - 1 <= dimension <
                num_dimensions_original + 1):
            raise ValueError(
                f"Dimension out of range. Expected `dimension` to be in range of "
                f" [-a.ndim - 1, a.ndim + 1), i.e. "
                f" [{-num_dimensions_original - 1}, {num_dimensions_original + 1}).")

        if dimension < 0:
            dimension = dimension + num_dimensions_original + 1

        original_shape: tuple = a.shape    # type: ignore[assignment]
        new_shape = list(original_shape)
        new_shape.insert(dimension, 1)
        return mi.TensorXf(a.array, shape=new_shape)