squeeze

nvtripy.squeeze(input: Tensor, dims: Sequence[int] | int) → Tensor

Returns a new tensor with the specified singleton dimensions of the input tensor removed.

Parameters:

• input (Tensor) – [dtype=T1] The input tensor.

• dims (Sequence[int] | int) – The dimension(s) to remove. These must have a length of 1.

Returns:

[dtype=T1] A new tensor.

Return type:

Tensor

DATA TYPE CONSTRAINTS:

• T1: bool, float16, bfloat16, float32, int8, int32, int64

Example: Squeeze All Dimensions

input = tp.iota((1, 2, 1), dtype=tp.float32)
output = tp.squeeze(input, dims=(0, 2))

Local Variables

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

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

Example: Squeeze First Dimension

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

Local Variables

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

>>> output
tensor(
    [[0],
     [0]], 
    dtype=float32, loc=gpu:0, shape=(2, 1))

Example: Squeeze First And Third Dimension

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

Local Variables

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

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