jax_sum

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

Compute the sum, mirroring jax.numpy.sum.

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.

  • initial – The starting value of the reduction.

  • where – Optional mask selecting the elements to include.

Returns:

The sum of the array.