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

DATA TYPE CONSTRAINTS:

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

>>> output
tensor([0, 0], 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, 0, 0, 0],
      [0, 0, 0, 0],
      [0, 0, 0, 0]],

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

>>> output
tensor(
    [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
      0, 0],
     [1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
      1, 1]], 
    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, 0, 0, 0, 0],
       [0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0]],

      [[0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0]],

      [[0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0]]],


     [[[1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1]],

      [[1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1]],

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

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

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