code: 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
  
imports: import contextlib
  import enum
  import gin
  from trax.fastmath.jax import JAX_BACKEND
  from trax.fastmath.numpy import NUMPY_BACKEND
  from trax.fastmath.tf import TF_BACKEND
label_desc: Use image node_id 33 for calling a global function with example usage: custom_vjp(f, f_fwd, f_bwd, nondiff_argnums) and returns: f_joint, backend, backend, f_diff_vjp, out,  vjpfn, f, out,  vjpfn, f_bwd, res