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","."]
docstr: Sums locally and then over devices with the axis name provided.
docstr_tok: ["Sums","locally","and","then","over","devices","with","the","axis","name","provided","."]
code: def __init__(self, axis_name: Optional[str] = None):
      """Sums locally and then over devices with the axis name provided."""
      self._axis_name = axis_name
  
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 1 to create a new AllSum object  with example: obj = AllSum(axis_name)