code: def inverse_fn(fn):
      """Compute the inverse of a function.
  
      This assumes that `fn` has a field called `inverse` which contains the inverse
      of the function.
  
      Args:
        fn: Function to invert.
  
      Returns:
        inverse: Inverse 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))
  
      def scale_by_half(x):
        return x / 2, (x, -jnp.log(2))
  
      scale_by_two.inverse = scale_by_half
      scale_by_half.inverse = scale_by_two
  
      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)
  
      inv_scale_by_2 = inverse_fn(scale_by_2)
      assert inv_scale_by_2 == scale_by_half
  
      x, x_extra, x_ldj = value_and_ldj(inv_scale_by_2, 4.)
      assert x == 2
      assert x_extra == 4
      assert x_ldj == -jnp.log(2)
      ```
      """
      return fn.inverse
  
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 17 for calling a global function with example usage: inverse_fn(fn) and returns: fn