mean¶

nvtripy.mean(input: Tensor, dim: int | Sequence[int] | None = None, keepdim: bool = False) → Tensor[source]¶

Returns a new tensor containing the mean of the elements of the input tensor along the specified dimension.

Parameters:

input (Tensor) – The input tensor.
dim (int | Sequence[int] | None) – The dimension or dimensions along which to reduce. If this is not provided, all dimensions are reduced.
keepdim (bool) – Whether to retain reduced dimensions in the output. If this is False, reduced dimensions will be squeezed.

Returns:

mean of the input tensor

Return type:

Tensor

INPUT REQUIREMENTS:: input.dtype is one of [float32, int32, int64, float16, bfloat16]
OUTPUT GUARANTEES:: return[0].dtype == input.dtype

Example

input = tp.reshape(tp.arange(6, dtype=tp.float32), (2, 3))
output = tp.mean(input, dim=1, keepdim=True)

Local Variables¶

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

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