jactorch.functional.shape

Tensor shape utilities.

Functions

add_dim(tensor, dim, size)	Add a dimension at dim with size size.
add_dim_as_except(tensor, target, *excepts)	Add dimension for the input tensor so that
broadcast(tensor, dim, size)	Broadcast a specific dim for size times.
broadcast_as_except(tensor, target, *excepts)	Add AND expand dimension for the input tensor so that
concat_shape(*shapes)	Concatenate shapes into a tuple.
flatten(tensor)	Flatten the tensor.
flatten2(tensor)	Flatten the tensor while keep the first (batch) dimension.
force_view(tensor, *shapes)	Do a view with optional contiguous copy.
move_dim(tensor, dim, dest)	Move a specific dimension to a designated dimension.
repeat(tensor, dim, count)	Repeat a specific dimension for count times.
repeat_times(tensor, dim, repeats)	Repeat each element along a specific dimension for repeats times.

Functions

add_dim(tensor, dim, size)

Add a dimension at dim with size size.

Example

>>> add_dim(torch.tensor([1, 2, 3]), 0, 2)
tensor([[1, 2, 3],
        [1, 2, 3]])

Parameters:

• tensor (Tensor) – the tensor to be broadcasted.

• dim (int) – the new dimension to be added.

• size (int) – the size of the target dimension.

Returns:

the broadcasted tensor.

Return type:

Tensor

add_dim_as_except(tensor, target, *excepts)

Add dimension for the input tensor so that

• It has the same number of dimensions as target.

• The original axes of the tensor are ordered in excepts.

Basically, it will “match” the shape of the target tensor, and align the axes in the input tensor with excepts in the target tensor.

Note

The list excepts must be in ascending order.

See also

broadcast_as_except(), which adds AND expands dimension.

Example

>>> add_dim_as_except(torch.rand(3, 4), torch.rand(2, 3, 4, 5), 1, 1).size()
torch.Size([1, 3, 4, 1])

Parameters:

• tensor (Tensor) – the tensor to be broadcasted.

• target (Tensor) – the target tensor.

• excepts (int) – the dimensions to be kept.

Returns:

the broadcasted tensor.

Return type:

Tensor

broadcast(tensor, dim, size)

Broadcast a specific dim for size times. Originally the dim size must be 1.

Example

>>> broadcast(torch.tensor([1, 2, 3]), 0, 2)
tensor([[1, 2, 3],
        [1, 2, 3]])

Parameters:

• tensor (Tensor) – the tensor to be broadcasted.

• dim (int) – the dimension to be broadcasted.

• size (int) – the size of the target dimension.

Returns:

the broadcasted tensor.

Return type:

Tensor

broadcast_as_except(tensor, target, *excepts)

Add AND expand dimension for the input tensor so that

• It has the same number of dimensions as target.

• The original axes of the tensor are ordered in excepts.

• The original axes of the tensor are expanded to the size of the corresponding axes in target.

After this function, the input tensor will have the same shape as the target tensor.

Note

The list excepts must be in ascending order.

See also

add_dim_as_except(), which only adds dimension (without expanding).

Example

>>> broadcast_as_except(torch.rand(3, 4), torch.rand(2, 3, 4, 5), 1, 1).size()
torch.Size([2, 3, 4, 5])

Parameters:

• tensor (Tensor) – the tensor to be broadcasted.

• target (Tensor) – the target tensor.

• excepts (int) – the dimensions to be kept.

Returns:

the broadcasted tensor.

Return type:

Tensor

concat_shape(*shapes)

Concatenate shapes into a tuple. The values can be either torch.Size, tuple, list, or int.

Parameters:

shapes (Size | Tuple[int, ...] | List[int] | int)

Return type:

Tuple[int, …]

flatten(tensor)

Flatten the tensor.

Parameters:

tensor (Tensor)

Return type:

Tensor

flatten2(tensor)

Flatten the tensor while keep the first (batch) dimension.

Parameters:

tensor (Tensor)

Return type:

Tensor

force_view(tensor, *shapes)

Do a view with optional contiguous copy. DEPRECATED. Use tensor.reshape instead.

Parameters:

• tensor (Tensor)

• shapes (int)

Return type:

Tensor

move_dim(tensor, dim, dest)

Move a specific dimension to a designated dimension.

Parameters:

• tensor (Tensor)

• dim (int)

• dest (int)

Return type:

Tensor

repeat(tensor, dim, count)

Repeat a specific dimension for count times.

Example

>>> repeat(torch.tensor([1, 2, 3]), 0, 2)
tensor([1, 2, 3, 1, 2, 3])

See also

repeat_times(), which repeats each element along the dimension.

Parameters:

• tensor (Tensor) – the tensor to be repeated.

• dim (int) – the dimension to be repeated.

• count (int) – the number of times to repeat.

Returns:

the repeated tensor.

Return type:

Tensor

repeat_times(tensor, dim, repeats)

Repeat each element along a specific dimension for repeats times.

Example

>>> repeat_times(torch.tensor([1, 2, 3]), 0, 2)
tensor([1, 1, 2, 2, 3, 3])

See also

repeat(), which repeats the whole dimension for repeats times.

Parameters:

• tensor (Tensor) – the tensor to be repeated.

• dim (int) – the dimension to be repeated.

• repeats (int) – the number of times to repeat each element.

Returns:

the repeated tensor.

Return type:

Tensor