insert_dims

sionna.phy.utils.insert_dims(tensor: torch.Tensor, num_dims: int, axis: int = -1) → torch.Tensor

Adds multiple length-one dimensions to a tensor.

This operation is an extension to PyTorch’s unsqueeze function. It inserts num_dims dimensions of length one starting from the dimension axis of a tensor. The dimension index follows Python indexing rules, i.e., zero-based, where a negative index is counted backward from the end.

Parameters:

• tensor (torch.Tensor) – Input tensor

• num_dims (int) – Number of dimensions to add

• axis (int) – Dimension index at which to expand the shape of tensor. Given a tensor of D dimensions, axis must be within the range [-(D+1), D] (inclusive).

Outputs:

tensor – A tensor with the same data as tensor, with num_dims additional dimensions inserted at the index specified by axis

Examples

import torch
from sionna.phy.utils import insert_dims

x = torch.ones([3, 4])
print(x.shape)
# torch.Size([3, 4])

y = insert_dims(x, num_dims=2, axis=-1)
print(y.shape)
# torch.Size([3, 4, 1, 1])