def custom_vjp(f, f_fwd, f_bwd, nondiff_argnums=()):
  """Set a custom vjp computation (override the default) for a function."""
  # Call backend custom_vjp if it exists.
  # TODO(lukaszkaiser): unify the APIs and remove nondiff_argnums altogether.
  if 'custom_vjp' in backend():
    return backend()['custom_vjp'](f, f_fwd, f_bwd)

  # Check that nondiff_argnums is (0, 1, ..., N) for some N.
  # Currently we only support nondiff_argnums at the front.
  counter = -1
  for i in nondiff_argnums:
    counter += 1
    if i != counter:
      raise ValueError('Currently we only support custom_vjps with all nondiff'
                       '_argnums up front, like (0,) or (0, 1) but not (1,) or'
                       ' (1, 2). Found: %s' % str(nondiff_argnums))

  # Use custom_grad.
  if counter == -1:  # no non-diff args
    def f_vjp(*args):
      out, residual = f_fwd(*args)
      def vjpfn(g):
        return f_bwd(residual, g)
      return out, vjpfn
    return backend()['custom_grad'](f_vjp, f)

  # Handle non-diff args by closure.
  def f_joint(*args):
    """This function takes all args, first counter+1 are non-diff ones."""
    nondiff_args = list(args[:counter+1])
    def f_diff(*diff_args):  # Takes only diff args, will define custom grad.
      args = nondiff_args + list(diff_args)
      return f(*args)
    def f_vjp(*diff_args):  # Custom VJP for diff args.
      args = nondiff_args + list(diff_args)
      out, residual = f_fwd(*args)
      def vjpfn(g):
        bwd_args = [residual, g]
        res = f_bwd(*bwd_args)
        return res[counter+1:]
      return out, vjpfn
    # This is the function taking only diff args with custom vjp.
    f_diff_vjp = backend()['custom_grad'](f_vjp, f_diff)
    # Call it on the diff args.
    return f_diff_vjp(*args[counter+1:])
  return f_joint
Set a custom vjp computation (override the default) for a function.
trax/fastmath/ops.py