squeeze¶

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

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

Parameters:

input (Tensor) – [dtype=T1] The input tensor.
dims (Sequence[int] | int) – The singleton dimension(s) to be removed. If this is not provided, all dimensions of size 1 are removed.

Raises:

TripyException – If any of the specified dimensions have a size that is not equal to 1.

Returns:

[dtype=T1] A new tensor.

Return type:

Tensor

TYPE CONSTRAINTS:

T1: bool, float8, bfloat16, float16, 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.0000],
      [0.0000]]], 
    dtype=float32, loc=gpu:0, shape=(1, 2, 1))

>>> output
tensor([0.0000, 0.0000], 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.0000],
      [0.0000]]], 
    dtype=float32, loc=gpu:0, shape=(1, 2, 1))

>>> output
tensor(
    [[0.0000],
     [0.0000]], 
    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.0000],
      [0.0000]]], 
    dtype=float32, loc=gpu:0, shape=(1, 2, 1))

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