code: def _sharded_log_prob_ratio(p, x, q, y, name=None):
      """Distributed log-prob ratio for Sharded."""
      with tf.name_scope(name or "sharded_log_prob_ratio"):
          if (
              p.experimental_shard_axis_names
              != q.experimental_shard_axis_names
          ):
              raise ValueError(
                  "Mismatched axis names "
                  f'"{p.experimental_shard_axis_names}" vs "'
                  f'"{q.experimental_shard_axis_names}"'
              )
  
          def log_prob_ratio_fn(x, y):
              return log_prob_ratio.log_prob_ratio(
                  p.distribution, x, q.distribution, y
              )
  
          axes = p.experimental_shard_axis_names
          return distribute_lib.make_psum_function(
              log_prob_ratio_fn,
              in_axes=(axes, axes),
              out_axes=axes,
              out_dtype=x,
          )(x, y)
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.experimental.bijectors import sharded
  from tensorflow_probability.python.internal import auto_composite_tensor
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import samplers
label_desc: Use image node_id 2 for calling a global function with example usage: _sharded_log_prob_ratio(p, x, q, y, name) and returns: distribute_lib, log_prob_ratio