unsqueeze¶

nvtripy.unsqueeze(input: Tensor, dim: int) → Tensor[source]¶

Returns a new tensor with the contents of the input tensor with a singleton dimension inserted before the specified axis.

Parameters:

input (Tensor) – [dtype=T1] The input tensor.
dim (int) – index before which to insert the singleton dimension. A negative dimension will be converted to dim = dim + input.rank + 1.

Returns:

[dtype=T1] A new tensor.

Return type:

Tensor

TYPE CONSTRAINTS:

T1: bool, float8, float16, bfloat16, float32, int4, int8, int32, int64

Example

input = tp.iota((2, 2), dtype=tp.float32)
output = tp.unsqueeze(input, 1)

Local Variables¶

>>> input
tensor(
    [[0.0000, 0.0000],
     [1.0000, 1.0000]], 
    dtype=float32, loc=gpu:0, shape=(2, 2))

>>> output
tensor(
    [[[0.0000, 0.0000]],

     [[1.0000, 1.0000]]], 
    dtype=float32, loc=gpu:0, shape=(2, 1, 2))