def transform_log_prob_fn(log_prob_fn: PotentialFn,
                          bijector: BijectorNest,
                          init_state: State = None
                         ) -> Union[PotentialFn, Tuple[PotentialFn, State]]:
  """
  """

  def wrapper(*args):
    """Transformed wrapper."""
    bijector_ = bijector

    args = tf.nest.map_structure(lambda x: 0. + x, args)
    if len(args) == 1:
      args = args[0]
    elif isinstance(bijector_, list):
      bijector_ = tuple(bijector_)

    original_space_args = tf.nest.map_structure(lambda b, x: b.forward(x),
                                                bijector_, args)
    original_space_args = original_space_args  # type: Tuple[Any]
    original_space_log_prob, extra = call_fn(log_prob_fn, original_space_args)
    event_ndims = tf.nest.map_structure(
        lambda x: tf.rank(x) - tf.rank(original_space_log_prob), args)

    return original_space_log_prob + sum(
        tf.nest.flatten(
            tf.nest.map_structure(
                lambda b, x, e: b.forward_log_det_jacobian(x, event_ndims=e),
                bijector_, args, event_ndims))), [original_space_args, extra]

  if init_state is None:
    return wrapper
  else:
    return wrapper, tf.nest.map_structure(lambda b, s: b.inverse(s), bijector,
                                          init_state)