concatenate

nvtripy.concatenate(tensors: Sequence[Tensor], dim: int) → Tensor[source]

Concatenates the input tensors along the specified dimension.

Parameters:

  • tensors (Sequence[Tensor]) – Sequence of tensors to concatenate. They must have identical shapes expect on the concatenation dimension.

  • dim (int) – The dimension along which the tensors are concatenated.

Returns:

Concatenated tensor whose shape is the same as the inputs except along dim, whose length is the sum of the lengths of the dim axis of the inputs.

Return type:

Tensor

INPUT REQUIREMENTS:

tensors.dtype is one of [float32, float16, bfloat16, float8, int4, int8, int32, int64, bool]

OUTPUT GUARANTEES:

return[0].dtype == tensors.dtype

Example
1a = tp.iota((1, 2), dtype=tp.float32)
2b = tp.iota((2, 2), dtype=tp.float32)
3
4output = tp.concatenate([a, b], dim=0)
Local Variables
>>> a
tensor(
    [[0, 0]], 
    dtype=float32, loc=gpu:0, shape=(1, 2))

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

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