code: def pmax(x, named_axis=None, allow_all_gather=False):
      """Generic `pmax` implementation."""
      # TODO(b/187173243): fix gradients for pmax
      axes = canonicalize_named_axis(named_axis)
      for axis in axes:
          if JAX_MODE:
              x = lax.pmax(x, axis)
          elif allow_all_gather:
              ctx = tf.distribute.get_replica_context()
              x = tf.reduce_max(
                  ctx.all_gather(x[tf.newaxis], axis=0), axis=0
              )
          else:
              raise NotImplementedError(
                  "`pmax` has no native implementation in TF. Pass in "
                  "`allow_all_gather=True` to enable a potentially "
                  "inefficient `all_gather`-based fallback. Also see b/191501877."
              )
      return x
  
imports: import collections
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow.python.util import nest
label_desc: Use image node_id 6 for calling a global function with example usage: pmax(x, named_axis, allow_all_gather) and returns: x