DimensionSize¶

class nvtripy.DimensionSize(data: int, name: str | None = None)[source]¶

Bases: Tensor

A 0D, int32 tensor that represents a scalar value extracted from the shape of a tensor.

Parameters:

data (int) – The value of the DimensionSize, which should be a scalar integer.
name (str | None) – An optional name.

eval() → DimensionSize[source]¶

Immediately evaluates this DimensionSize object.

Note

DimensionSize will always reside on host even after it is evaluated.

Returns:: The evaluated DimensionSize.
Return type:: DimensionSize

Example

dim_size = tp.ones((2, 2)).shape[0]
dim_size.eval()
print(dim_size.device)

Local Variables¶

>>> dim_size
2

Output¶

cpu:0

cast(dtype: dtype) → Tensor¶

Returns a tensor with the contents of the input tensor casted to the specified data type.

For casts into quantized datatypes (int4 and float8), this performs a per-tensor quantization into that datatype with scale 1.0; for casts from those datatypes, this performs a per-tensor dequantization with scale 1.0. Direct use of quantize() and dequantize() allows for finer control over these parameters.

Parameters:

input (Tensor) – [dtype=T1] The input tensor.
dtype (dtype) – [dtype=T2] The desired data type.

Returns:

[dtype=T2] A tensor containing the casted values.

Return type:

Tensor

DATA TYPE CONSTRAINTS:

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

UNSUPPORTED DATA TYPE COMBINATIONS:

T1=float8, T2=int4
T1=float8, T2=int8
T1=int8, T2=float8
T1=int4, T2=float8
T1=int4, T2=int8
T1=int4, T2=int64

Example

input = tp.Tensor([1, 2])
output = tp.cast(input, tp.float32)

Local Variables¶

>>> input
tensor([1, 2], dtype=int32, loc=cpu:0, shape=(2,))

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

See also

quantize(), dequantize()

copy(device: device) → Tensor¶

Copies the input tensor to the specified device.

Caution

This function cannot be used in a compiled function or nvtripy.Module because it depends on evaluating its inputs, which is not allowed during compilation.

Parameters:

input (Tensor) – [dtype=T1] Input tensor.
device (device) – The target device to copy the tensor to.

Returns:

[dtype=T1] A new tensor on the specified device.

Raises:

TripyException – If the input tensor is already on the specified device, as performing copies within the same device is currently not supported.

Return type:

Tensor

DATA TYPE CONSTRAINTS:

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

Example: Copying To CPU

input = tp.Tensor([1, 2, 3], device=tp.device("gpu"))
output = tp.copy(input, device=tp.device("cpu"))

Local Variables¶

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

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

Example: Copying To GPU

input = tp.Tensor([1, 2, 3])
output = tp.copy(input, device=tp.device("gpu"))

Local Variables¶

>>> input
tensor([1, 2, 3], dtype=int32, loc=cpu:0, shape=(3,))

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

flatten(start_dim: int = 0, end_dim: int = -1) → Tensor¶

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))

permute(perm: Sequence[int]) → Tensor¶

Returns a tensor with its dimensions permuted.

Parameters:

input (Tensor) – [dtype=T1] The input tensor.
perm (Sequence[int]) – The desired ordering of dimensions. It must contain all integers in \([0..N-1]\) exactly once, where \(N\) is the rank of the input tensor.

Returns:

[dtype=T1] A new tensor.

Return type:

Tensor

DATA TYPE CONSTRAINTS:

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

Example

input = tp.reshape(tp.arange(6, dtype=tp.float32), (2, 3))
output = tp.permute(input, (1, 0))

Local Variables¶

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

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

reshape(shape: Sequence[int | DimensionSize]) → Tensor¶

Returns a new tensor with the contents of the input tensor in the specified shape.

Parameters:

input (Tensor) – [dtype=T1] The input tensor.
shape (Sequence[int | DimensionSize]) – The desired compatible shape. If a shape dimension is -1, its value is inferred based on the other dimensions and the number of elements in the input. Atmost one dimension can be -1.

Returns:

[dtype=T1] A new tensor with the specified shape.

Return type:

Tensor

DATA TYPE CONSTRAINTS:

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

Example

input = tp.iota((2, 3), dtype=tp.float32)
output = tp.reshape(input, (1, 6))

Local Variables¶

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

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

squeeze(dims: Sequence[int] | int) → Tensor¶

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

Parameters:

input (Tensor) – [dtype=T1] The input tensor.
dims (Sequence[int] | int) – The dimension(s) to remove. These must have a length of 1.

Returns:

[dtype=T1] A new tensor.

Return type:

Tensor

DATA TYPE CONSTRAINTS:

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

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

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

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

tolist() → List | Number¶

Returns the tensor as a nested list. If the tensor is a scalar, returns a python number.

Returns:: The tensor represented as a nested list or a python number.
Return type:: List | Number

Example: Ranked tensor

tensor = tp.ones((2, 2))
tensor_list = tensor.tolist()

Local Variables¶

>>> tensor_list
[[1.0, 1.0], [1.0, 1.0]]

Example: Scalar

tensor = tp.Tensor(2.0)
tensor_scalar = tensor.tolist()

Local Variables¶

>>> tensor_scalar
2.0

transpose(dim0: int, dim1: int) → Tensor¶

Returns a new tensor that is a transposed version of the input tensor where dim0 and dim1 are swapped.

Parameters:

input (Tensor) – [dtype=T1] The input tensor.
dim0 (int) – The first dimension to be transposed.
dim1 (int) – The second dimension to be transposed.

Returns:

[dtype=T1] A new tensor.

Return type:

Tensor

DATA TYPE CONSTRAINTS:

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

Example

input = tp.reshape(tp.arange(6, dtype=tp.float32), (2, 3))
output = tp.transpose(input, 0, 1)

Local Variables¶

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

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

unsqueeze(dim: int) → Tensor¶

Returns a new tensor with the contents of the input tensor with a singleton dimension inserted before the specified axis.

Parameters:

input (Tensor) – [dtype=T1] The input tensor.
dim (int) – index before which to insert the singleton dimension. A negative dimension will be converted to dim = dim + input.rank + 1.

Returns:

[dtype=T1] A new tensor.

Return type:

Tensor

DATA TYPE CONSTRAINTS:

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

Example

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

Local Variables¶

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

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

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