code: def value_and_ldj(fn, args):
      """Compute the value and log-det jacobian of function evaluated at args.
  
      This assumes that `fn`'s `extra` output is a 2-tuple, where the first element
      is arbitrary and the the last element is the log determinant of the jacobian
      of the transformation.
  
      Args:
        fn: Function to evaluate.
        args: Arguments to `fn`.
  
      Returns:
        ret: First output of `fn`.
        extra: Second output of `fn`.
        ldj: Log-det jacobian of `fn`.
  
      #### Example
  
      ```python
      def scale_by_two(x):
        # Return x unchanged as the extra output for illustrative purposes.
        return 2 * x, (x, jnp.log(2))
  
      y, y_extra, y_ldj = value_and_ldj(scale_by_2, 3.)
      assert y == 6
      assert y_extra == 3
      assert y_ldj == jnp.log(2)
      ```
      """
      value, (extra, ldj) = fn(args)
      return value, (extra, ldj), ldj
  
imports: import functools
  import jax
  from jax import lax
  from jax import random
  from jax import tree_util
  import jax.numpy
label_desc: Use image node_id 16 for calling a global function with example usage: value_and_ldj(fn, args) and returns: value,  extra,  ldj,  ldj