flatten¶

nvtripy.flatten(input: Tensor, start_dim: int = 0, end_dim: int = -1) → Tensor[source]¶

Flattens the input tensor from start_dim to end_dim.

Parameters:

input (Tensor) – [dtype=T1] The input tensor to be flattened.
start_dim (int) – The first dimension to flatten (default is 0).
end_dim (int) – The last dimension to flatten (default is -1, which includes the last dimension).

Returns:

[dtype=T1] A flattened tensor.

Return type:

Tensor

TYPE CONSTRAINTS:

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

Example: Flatten All Dimensions

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

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: Flatten Starting from First Dimension

input = tp.iota((2, 3, 4), dtype=tp.float32)
output = tp.flatten(input, start_dim=1)

Local Variables¶

>>> input
tensor(
    [[[0.0000, 0.0000, 0.0000, 0.0000],
      [0.0000, 0.0000, 0.0000, 0.0000],
      [0.0000, 0.0000, 0.0000, 0.0000]],

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

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

Example: Flatten a Specific Range of Dimensions

input = tp.iota((2, 3, 4, 5), dtype=tp.float32)
output = tp.flatten(input, start_dim=1, end_dim=2)

Local Variables¶

>>> input
tensor(
    [[[[0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
       [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
       [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
       [0.0000, 0.0000, 0.0000, 0.0000, 0.0000]],

      [[0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
       [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
       [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
       [0.0000, 0.0000, 0.0000, 0.0000, 0.0000]],

      [[0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
       [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
       [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
       [0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]],


     [[[1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
       [1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
       [1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
       [1.0000, 1.0000, 1.0000, 1.0000, 1.0000]],

      [[1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
       [1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
       [1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
       [1.0000, 1.0000, 1.0000, 1.0000, 1.0000]],

      [[1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
       [1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
       [1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
       [1.0000, 1.0000, 1.0000, 1.0000, 1.0000]]]], 
    dtype=float32, loc=gpu:0, shape=(2, 3, 4, 5))

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

     [[1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
      [1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
      [1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
      ...,
      [1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
      [1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
      [1.0000, 1.0000, 1.0000, 1.0000, 1.0000]]], 
    dtype=float32, loc=gpu:0, shape=(2, 12, 5))