code: def compute_and_maybe_broadcast_ldj(
      b, x, event_ndims, ldj_fn=lambda b: b.forward_log_det_jacobian
  ):
      """Broadcasts the forward/inverse log det jacobian to full batch shape."""
      ldj = ldj_fn(b)(x, event_ndims=event_ndims)
      x_batch_shape_parts = [
          ps.shape(t)[: ps.rank(t) - nd]
          for (t, nd) in zip(
              tf.nest.flatten(x), tf.nest.flatten(event_ndims)
          )
      ]
      return tf.broadcast_to(
          ldj,
          functools.reduce(
              ps.broadcast_shape, x_batch_shape_parts, ps.shape(ldj)
          ),
      )
  
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import generic
label_desc: Use image node_id 5 for calling a global function with example usage: compute_and_maybe_broadcast_ldj(b, x, event_ndims, ldj_fn) and returns: tf