jax_mean

probly.representation.jax_functions.jax_mean(a: JaxArrayLike, /, axis: int | Sequence[int] | None = None, dtype: DTypeLike | None = None, out: None = None, keepdims: bool = False, *, where: JaxArrayLike | None = None) → jax.Array[source]

Compute the arithmetic mean, mirroring jax.numpy.mean.

Parameters:

  • a – The array to reduce.

  • axis – The axis or axes to reduce.

  • dtype – The data type of the accumulator.

  • out – Unsupported by JAX, must be None.

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

  • where – Optional mask selecting the elements to include.

Returns:

The mean of the array.