probly.representation.torch_functions.torch_average

probly.representation.torch_functions.torch_average(tensor: Tensor, /, dim: int | tuple[int, ...] | None = None, *, axis: int | tuple[int, ...] | None = None, weights: Tensor | None = None, keepdim: bool = False, out: Tensor | None = None) → Tensor[source]

Compute a possibly weighted average over tensor dimensions.

Parameters:

  • tensor – Tensor to average.

  • dim – Dimension or dimensions to reduce. If omitted, all dimensions are reduced.

  • axis – Alias for dim for NumPy-style call sites.

  • weights – Optional weights. Weights may have the same shape as input, be broadcastable to input, or match the reduced dimensions.

  • keepdim – Whether reduced dimensions are retained with size one.

  • out – Optional output tensor.

Returns:

The weighted average.