expand_to_rank

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

Inserts as many axes to a tensor as needed to achieve a desired rank.

This operation inserts additional dimensions to a tensor starting at axis, so that the rank of the resulting tensor has rank target_rank. 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

• target_rank (int) – Rank of the output tensor. If target_rank is smaller than the rank of tensor, the function does nothing.

• 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 target_rank - rank(tensor) additional dimensions inserted at the index specified by axis. If target_rank <= rank(tensor), tensor is returned.

Examples

import torch
from sionna.phy.utils import expand_to_rank

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

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