code: def _cross_replica_scatter_add(
      source: Array, indices: Array, updates: Array, axis_name
  ):
      """tf.scatter_add, but with JAX, cross replica, and without state.
  
      Args:
        source: An array of shape [O].
        indices: An array indicating which index each update is for.
        updates: The updates to apply to `source`. Of same shape as indices.
        axis_name: What axis to aggregate over, if str. If passed an iterable,
          aggregates over multiple axes. Defaults to no aggregation, i.e. None.
  
      Returns:
        An array of shape [O], which is source + the scattered updates from all
        replicas.
      """
      assert updates.shape == indices.shape
      assert jnp.issubdtype(indices.dtype, jnp.integer)
      assert source.ndim == 1
      # Flatten indices, updates.
      num_classes = source.shape[0]
      indices = jnp.reshape(indices, [-1])
      updates = jnp.reshape(updates, [-1])
      # Scatter updates according to value of indices.
      updates_at_idxs = updates[..., None] * base.one_hot(
          indices, num_classes
      )
      # Aggregate locally first, then across replicas.
      total_updates = jnp.sum(updates_at_idxs, axis=0)
      if axis_name is not None:
          axis_names = (
              (axis_name,) if isinstance(axis_name, str) else axis_name
          )
          for a_name in axis_names:
              total_updates = jax.lax.psum(
                  total_updates, axis_name=a_name
              )
      return source + total_updates
  
file_path: rlax/rlax/_src/pop_art.py
imports: import collections
  from typing import Mapping, Tuple
  import chex
  import jax
  import jax.numpy
  from rlax._src import base
label_desc: Use image node_id 1 for calling a global function with example usage: _cross_replica_scatter_add(source, indices, updates, axis_name) and returns: unknown