code: def psum_metric_normalizer(
      metrics: Tuple[jnp.ndarray, jnp.ndarray]
  ) -> Tuple[jnp.ndarray, jnp.ndarray]:
      """Applies psum over the given tuple of (metric, normalizer)."""
      psumed_metric = jnp.sum(
          jax.lax.psum(metrics[0], axis_name="batch")
      )
      psumed_normalizer = jnp.sum(
          jax.lax.psum(metrics[1], axis_name="batch")
      )
      return (psumed_metric, psumed_normalizer)
  
imports: from typing import Tuple, Optional
  import jax
  import jax.numpy
label_desc: Use image node_id 4 for calling a global function with example usage: psum_metric_normalizer(metrics) and returns: psumed_metric,  psumed_normalizer