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