code: def _axis_size(x, axis=None):
      """Get number of elements of `x` in `axis`, as type `x.dtype`."""
      if axis is None:
          return ps.cast(ps.size(x), x.dtype)
      return ps.cast(
          ps.reduce_prod(ps.gather(ps.shape(x), axis)), x.dtype
      )
  
file_path: probability/tensorflow_probability/python/mcmc/diagnostic.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python import stats
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow.python.util import nest
label_desc: Use image node_id 6 for calling a global function with example usage: _axis_size(x, axis) and returns: ps, ps