code: def make_transformed_log_prob(
      log_prob_fn, bijector, direction, enable_bijector_caching=True
  ):
      """Transforms a log_prob function using bijectors.
  
      Note: `direction = 'inverse'` corresponds to the transformation calculation
      done in `tfp.distributions.TransformedDistribution.log_prob`.
  
      Args:
        log_prob_fn: Python `callable` taking an argument for each state part which
          returns a `Tensor` representing the joint `log` probability of those state
          parts.
        bijector: `tfp.bijectors.Bijector`-like instance (or list thereof)
          corresponding to each state part. When `direction = 'forward'` the
          `Bijector`-like instance must possess members `forward` and
          `forward_log_det_jacobian` (and corresponding when
          `direction = 'inverse'`).
        direction: Python `str` being either `'forward'` or `'inverse'` which
          indicates the nature of the bijector transformation applied to each state
          part.
        enable_bijector_caching: Python `bool` indicating if `Bijector` caching
          should be invalidated.
          Default value: `True`.
  
      Returns:
        transformed_log_prob_fn: Python `callable` which takes an argument for each
          transformed state part and returns a `Tensor` representing the joint `log`
          probability of the transformed state parts.
      """
      if direction not in {"forward", "inverse"}:
          raise ValueError(
              'Argument `direction` must be either `"forward"` or '
              '`"inverse"`; saw "{}".'.format(direction)
          )
      fn = make_transform_fn(bijector, direction)
      ldj_fn = make_log_det_jacobian_fn(bijector, direction)
  
      def transformed_log_prob_fn(*state_parts):
          """Log prob of the transformed state."""
          if not enable_bijector_caching:
              state_parts = [tf.identity(sp) for sp in state_parts]
          tlp = log_prob_fn(*fn(state_parts))
          tlp_rank = prefer_static.rank(tlp)
          event_ndims = [
              (prefer_static.rank(sp) - tlp_rank) for sp in state_parts
          ]
          return tlp + ldj_fn(state_parts, event_ndims)
  
      return transformed_log_prob_fn
  
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow.python.util import deprecation
label_desc: Use image node_id 3 for calling a global function with example usage: make_transformed_log_prob(log_prob_fn, bijector, direction, enable_bijector_caching) and returns: transformed_log_prob_fn, unknown