code: def get_axis_size(axis_name=None):
      if JAX_MODE:
          return lax.psum(1, axis_name)
      ctx = tf.distribute.get_replica_context()
      return ctx.num_replicas_in_sync
  
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 10 for calling a global function with example usage: get_axis_size(axis_name) and returns: ctx, lax