jax_average

probly.representation.jax_functions.jax_average(a: JaxArrayLike, /, axis: int | Sequence[int] | None = None, weights: JaxArrayLike | None = None, returned: bool = False, keepdims: bool = False) → jax.Array | tuple[jax.Array, jax.Array][source]

Compute a possibly weighted average, mirroring jax.numpy.average.

Parameters:

  • a – The array to reduce.

  • axis – The axis or axes to reduce.

  • weights – Optional weights, broadcastable to a or matching the reduced axis.

  • returned – Whether to also return the sum of the weights.

  • keepdims – Whether reduced axes are retained with size one.

Returns:

The weighted average, and the sum of the weights if returned is True.