class_docstr: Helper for summing over elements in an array and over devices.
class_docstr_tok: ["Helper","for","summing","over","elements","in","an","array","and","over","devices","."]
code: def __call__(
      self, x: Array, axis: Optional[Union[int, Sequence[int]]] = None
  ) -> Array:
      s = jnp.sum(x, axis=axis)
      if self._axis_name:
          s = jax.lax.psum(s, axis_name=self._axis_name)
      return s
  
file_path: rlax/rlax/_src/base.py
imports: from typing import Optional, Sequence, Union
  import chex
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling the AllSum obj's underlying member method code with example usage: obj.__call__(x, axis) and returns: s