stack¶

tripy.stack(tensors: Sequence[Tensor], dim: int = 0) → Tensor[source]¶

Stacks multiple tensors of same shape along a given dimension.

Parameters:

tensors (Sequence[Tensor]) – [dtype=T1] Sequence of tensors of the same shape.
dim (int) – The dimension to insert.

Returns:

[dtype=T1] A tensor with a new dimension inserted at the specified position.

Return type:

Tensor

TYPE CONSTRAINTS:

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

Example

Example¶

a = tp.iota((2, 3), dtype=tp.float32)
b = tp.iota((2, 3), dtype=tp.float32)

output = tp.stack([a, b], dim=0)

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

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