code: def _make_gradient_fn_of_y(fn, x):
      """Defines the gradient of `fn(x, *args)` as a function of `(y, *args)`.
  
      The returned function `grad_fn_of_y(y, *args)` computes the partial
      derivative(s) of the invertible function `fn(x, *args)` with respect to `x`
      and any other arguments `*args`. It is conceptually defined by
  
      ```
      def grad_fn_of_y(y, *args):
        x = inverse_of_fn(y, *args)
        return gradient(fn, x, *args)
      ```
  
      Unfortunately, we don't have direct access to `inverse_of_fn`, so instead we
      require that the appropriate `x` is given to us, and the function we return
      is valid *only* at `y = fn(x, *args)` for the provided value of `x`.
  
      Why would we bother to define a function that's valid only at a single point?
      The real goal is to annotate the *derivatives* of `grad_fn_of_y`
      (which are second-order derivatives of `fn`) so that it appears to the
      autodiff system *as if* we had written the 'conceptual' code above, using a
      differentiable inverse function. Without this trick, we could still use
      a first-order derivative of `fn` to, e.g., define a bijector
      log-det-jacobian, but we wouldn't be able to differentiate it.
  
      Args:
        fn: Python `callable` scalar function of a scalar `x` (and
          optionally, of additional scalar arguments `*args`), twice differentiable
          and invertible in its first argument `x`.
        x: Float `Tensor` input at which `fn` and its derivatives are evaluated.
      Returns:
        grad_fn_of_y: Python `callable` that takes arguments `(y, *args)` and
          returns a list of first-order partial derivatives of `y = fn(x, *args)`.
          If `fn` takes only one argument, the value returned is a
          single-element list `[dy_dx]`. The value of `y` is assumed to be
          `y = fn(x, *args)` and is ignored, but derivatives wrt `y` and the other
          `*args` are defined.
  
      """
  
      def _grad_fn(x, *args):
          _, grads = gradient.value_and_gradient(fn, x, *args)
          return grads if args else [grads]  # Always return a list.
  
      def _grad_fn_of_y_fwd(unused_y, *args):
          dy_dx, *dy_dargs = _grad_fn(x, *args)
          return [
              (dy_dx, *dy_dargs),  # Actual return values.
              (
                  dy_dx,
                  dy_dargs,
                  args,
              ),  # Auxiliary values for the backward pass.
          ]
  
      def _second_order_terms(*args):
          """Computes entries of the (Hessian of `fn`) == (Jacobian of `_grad_fn`)."""
          # Partial derivatives of _grad_fn's first output (dy/dx) wrt `(x, *args)`.
          _, (d2y_dx2, *d2y_dx_dargs) = gradient.value_and_gradient(
              lambda x_and_args: _grad_fn(*x_and_args)[0],
              (x,) + args,
              auto_unpack_single_arg=False,
          )
  
          # Partial derivatives of additional outputs (dy/da, etc) wrt the input
          # *args (if any). Note that we don't need derivatives of these outputs wrt
          # `x`, since these are equal to the values we computed above in
          # `d2y_dx_dargs`by the [symmetry of partial derivatives](
          #   https://en.wikipedia.org/wiki/Symmetry_of_second_derivatives). This
          # could also in principle be applied to optimize redundant partial
          # derivatives computed in this loop, although this would be incompatible
          # with parallelizing the loop (which is probably a bigger win).
          d2y_dargs2 = []
          for i in range(len(args)):
              # It may be possible to run this loop in parallel with `vectorized_map`,
              # although this would only matter in cases with >> 1 arguments.
              _, d2y_dargs2_row = gradient.value_and_gradient(
                  lambda args, i=i: _grad_fn(x, *args)[1 + i],
                  args,
                  auto_unpack_single_arg=False,
              )
              d2y_dargs2.append(d2y_dargs2_row)
  
          return d2y_dx2, d2y_dx_dargs, d2y_dargs2
  
      def _jacobian_of_grad_fn_wrt_y(
          dy_dx, dy_dargs, d2y_dx2, d2y_dx_dargs, d2y_dargs2
      ):
          """Builds a Jacobian matrix for `_grad_fn_of_y(y, *args)`.
  
          Args:
            dy_dx: Tensor.
            dy_dargs: List of Tensors of length `len(args)`.
            d2y_dx2: Tensor.
            d2y_dx_dargs: List of Tensors of length `len(args)`.
            d2y_dargs2: Nested list of length `len(args)`, containing lists of length
              `len(args)`. Due to the symmetry of partial derivatives, the order of
              indices doesn't matter.
          Returns:
            jacobian: Nested list of length `len(args + 1)`, corresponding to the
              outputs of `_grad_fn_of_y` including the first output `dy/dx`. These
              contain inner lists, each also of `len(args + 1)`, corresponding
              to the inputs of `_grad_fn_of_y` including the first input `y`.
              The inner lists contain Tensor partial derivatives.
  
          ### Mathematical details
  
          Here we derive partial derivatives of the outputs of `grad_fn_of_y`:
  
          ```
          dy/dx, dy/da, dy/db, ... = grad_fn_of_y(y, a, b, ...)
                                   = grad_fn(x(y, a, b, ...), a, b, ...)
          ```
  
          with respect to its inputs. Our goal will be to express these partial
          derivatives, which are second-order derivatives of the original `fn`, in
          terms of quantities that we can evaluate by differentiating the forward
          `y = fn(x, a, b, ...)`; that is, we always want `dy` in the numerator.
  
          We can divide these into two cases (1) and (2) in the Jacobian matrix:
  
          ```
                                 inputs
                             y     a      b     ...
            outputs: dy/dx  (1)   (2)    (2)
                     dy/da  (1)   (2)    (2)   ...
                     dy/db  (1)   (2)    (2)
                     ...
          ```
  
          The first case (1) contains the derivatives with respect to `y`, which fall
          out straightforwardly from the inverse function theorem
          `dx/dy = 1 / (dy/dx)`, and the chain rule:
  
          ```
          d (dy/dx) / dy = d (dy/dx) / dx  * dx/dy
                         = d2y/dx2 /  dy/dx
          d (dy/da) / dy = d (dy/da) / dx  * dx/dy
                         = d2y/dxa /  dy/dx
          # etc.
          ```
  
          The second case (2) includes derivatives with respect to the additional
          arguments `a, b, ...`. These are more involved, because these arguments
          enter `grad_fn_of_y` twice: in addition to their direct effect on
          `grad_fn(x, a, b, ...)`, we must also account for the indirect effect via
          `x = inverse_of_fn(y, a, b, ...)`. This requires total differentiation:
  
          ```
          d f(x(y, a, b, ...), a, b, ...) / db = partial(f, x) * partial(x, b)
                                                + partial(f, a) * partial(a, b)
                                                + partial(f, b) * partial(b, b)
                                                + ...
          ```
  
          where the cross terms between constants, like `partial(a, b)`, are zero and
          disappear. Considering a concrete Jacobian entry `d (dy/da) / db`, we find:
  
          ```
          d (dy/da) / db = d2y/dab + d2y/dxa * dx/db
                         = d2y/dab - d2y/dxa * (dx/dy) * dy/db
                           # by the implicit differentiation identity
                           # `dx/db = -(dx/dy) * (dy/db)` (eqn (6) of Figurnov et al.
                         = d2y/dab - (d2y/dxa * dy/db) / (dy/dx)
                           # by the inverse function theorem.
          ```
  
          Analagous derivations, using any output in place of `dy/da`,
          and any non-`y` input in place of `b`, are sufficient to compute all of the
          remaining Jacobian terms.
  
          ```
  
          """
          # Avoid errors if any second derivatives of `fn` are `None`. It is assumed
          # that all first derivatives are non-`None`.
          zero_if_none = lambda v: 0.0 if v is None else v
  
          # Case (1): derivative of first output wrt first input.
          jacobian_first_row = [zero_if_none(d2y_dx2) / dy_dx]
          for d2y_dx_darg, dy_darg in zip(d2y_dx_dargs, dy_dargs):
              # Case (2): derivative of first output wrt subsequent input.
              jacobian_first_row.append(
                  zero_if_none(d2y_dx_darg)
                  - zero_if_none(d2y_dx2) * dy_darg / dy_dx
              )
  
          jacobian_all_rows = [jacobian_first_row]
          for d2y_dx_darg, d2y_dargs2_row in zip(
              d2y_dx_dargs, d2y_dargs2
          ):
              # Case (1): derivative of secondary return value wrt `y`.
              jacobian_current_row = [zero_if_none(d2y_dx_darg) / dy_dx]
              for arg_idx in range(len(dy_dargs)):
                  # Case (2): derivative of a secondary return value wrt secondary input.
                  jacobian_current_row.append(
                      zero_if_none(d2y_dargs2_row[arg_idx])
                      - zero_if_none(d2y_dx_darg)
                      * dy_dargs[arg_idx]
                      / dy_dx
                  )
              jacobian_all_rows.append(jacobian_current_row)
  
          return jacobian_all_rows
  
      def _grad_fn_of_y_bwd(aux, dresult_d_grad_fn):
          dy_dx, dy_dargs, args = aux
  
          jacobian = _jacobian_of_grad_fn_wrt_y(
              dy_dx, dy_dargs, *_second_order_terms(*args)
          )
  
          # We could concatenate the Jacobian components and do a real
          # `tf.linalg.matvec`, instead of coding a vector-matrix product manually,
          # but this is likely not a meaningful optimization in typical cases with
          # 1-2 args.
          gradients = []
          for j in range(1 + len(args)):
              g = 0.0
              for i in range(1 + len(args)):
                  g += dresult_d_grad_fn[i] * jacobian[i][j]
              gradients.append(g)
          return gradients
  
      def _grad_fn_of_y_jvp(primals, tangents):
          unused_y, *args = primals
          dy_dx, *dy_dargs = _grad_fn(x, *args)
          jacobian = _jacobian_of_grad_fn_wrt_y(
              dy_dx, dy_dargs, *_second_order_terms(*args)
          )
  
          # We could concatenate the Jacobian components and do a real
          # `tf.linalg.matvec`, instead of coding a matrix-vector product manually,
          # but this is likely not a meaningful optimization in typical cases with
          # 1-2 args.
          tangents_out = []
          for i in range(len(args) + 1):
              tangent_out = 0.0
              for j in range(len(args) + 1):
                  tangent_out += tangents[j] * jacobian[i][j]
              tangents_out.append(tangent_out)
  
          return (dy_dx, *dy_dargs), tuple(tangents_out)
  
      @tfp_custom_gradient.custom_gradient(
          vjp_fwd=_grad_fn_of_y_fwd,
          vjp_bwd=_grad_fn_of_y_bwd,
          jvp_fn=_grad_fn_of_y_jvp,
      )
      def _grad_fn_of_y_with_gradient(y, *args):
          del y  # Unused.
          return _grad_fn(x, *args)
  
      return _grad_fn_of_y_with_gradient
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.internal import callable_util
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math import root_search
label_desc: Use image node_id 1 for calling a global function with example usage: _make_gradient_fn_of_y(fn, x) and returns: _grad_fn_of_y_with_gradient, unknown, list, d2y_dx2,  d2y_dx_dargs,  d2y_dargs2, jacobian_all_rows, gradients, dy_dx,  unknown,  tuple, _grad_fn