class_docstr: Bijector to associate a numeric inverse with any invertible function.
class_docstr_tok: ["Bijector","to","associate","a","numeric","inverse","with","any","invertible","function","."]
docstr: Wraps the inverse to provide implicit reparameterization gradients.
docstr_tok: ["Wraps","the","inverse","to","provide","implicit","reparameterization","gradients","."]
code: def _wrap_inverse_with_implicit_gradient(self):
      """Wraps the inverse to provide implicit reparameterization gradients."""
  
      def _vjp_fwd(y, *args_with_full_batch_shape):
          # Prevent autodiff from trying to backprop through the root search.
          x = tf.stop_gradient(self._inverse_no_gradient(y))
          return x, (
              # Auxiliary values for the backwards pass.
              x,
              y,
              args_with_full_batch_shape,
          )
  
      # By the inverse function theorem, the derivative of an
      # inverse function is the reciprocal of the forward derivative. This has
      # been popularized in machine learning by [1].
      # [1] Michael Figurnov, Shakir Mohamed, Andriy Mnih (2018). Implicit
      #     Reparameterization Gradients. https://arxiv.org/abs/1805.08498.
      def _vjp_bwd(aux, dresult_d_x):
          x, y, args_with_full_batch_shape = aux
          dy_dx, *dy_dargs = _make_gradient_fn_of_y(self.fn, x)(
              y, *args_with_full_batch_shape
          )
          return [
              # Derivative of `y = f_inv(x)`.
              dresult_d_x
              / dy_dx
          ] + [
              # Derivatives to parameters by implicit differentiation. This
              # is a special case of eqn (6) in Figurnov et al. [1].
              -dresult_d_x * dy_darg / dy_dx
              for dy_darg in dy_dargs
          ]
  
      def _inverse_jvp(primals, tangents):
          y, *args_with_full_batch_shape = primals
          dy, *dargs = tangents
  
          # Prevent autodiff from trying to backprop through the root search.
          x = tf.stop_gradient(self._inverse_no_gradient(y))
  
          dy_dx, *dy_dargs = _make_gradient_fn_of_y(self.fn, x)(
              y, *args_with_full_batch_shape
          )
          dx = dy / dy_dx
          for darg, dy_darg in zip(dargs, dy_dargs):
              # Derivatives to parameters by implicit differentiation. This
              # is a special case of eqn (6) in Figurnov et al. [1].
              dx -= darg * dy_darg / dy_dx
          return x, dx
  
      @tfp_custom_gradient.custom_gradient(
          vjp_fwd=_vjp_fwd, vjp_bwd=_vjp_bwd, jvp_fn=_inverse_jvp
      )
      def _inverse_with_gradient(y, *args_with_full_batch_shape):
          del args_with_full_batch_shape  # Used by the backwards pass only.
          return self._inverse_no_gradient(y)
  
      def _arg_broadcasting_wrapped_inverse(y):
          args = self._additional_scalar_parameters_requiring_gradients
          if args:
              # Gradients of inverse wrt args will have full batch shape. If any arg
              # has a smaller shape, we need to deduplicate its gradient to have
              # that same smaller shape. An easy way to do that is to backprop
              # through `broadcast_to`. Note that this needs to occur outside of the
              # custom_gradient machinery, so that the backprop actually happens.
              # TODO(davmre): Do gradient reductions directly in the VJP using
              # `tf.raw_ops.BroadcastGradientArgs` so we can remove this wrapper
              # and avoid spurious broadcasting.
              full_batch_shape = ps.broadcast_shape(
                  self.experimental_batch_shape_tensor(), ps.shape(y)
              )
              args = [
                  tf.broadcast_to(arg, full_batch_shape) for arg in args
              ]
          return _inverse_with_gradient(y, *args)
  
      return _arg_broadcasting_wrapped_inverse
  
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 12 for calling the ScalarFunctionWithInferredInverse obj's underlying member method code with example usage: obj._wrap_inverse_with_implicit_gradient() and returns: _arg_broadcasting_wrapped_inverse, x,  x,  y,  args_with_full_batch_shape, unknown, x,  dx, self, _inverse_with_gradient