code: def _custom_vjp(f, f_fwd, f_bwd, nondiff_argnums=()):
      @functools.partial(
          jax.custom_vjp, nondiff_argnums=nondiff_argnums
      )
      def _f(*args, **kwargs):
          return f(*args, **kwargs)
  
      _f.defvjp(f_fwd, f_bwd)
      return _f
  
imports: import functools
  import jax
  from jax import lax
  from jax import random
  import jax.numpy
  import jax.scipy.special
  import numpy
  import tensorflow
  import tensorflow_datasets
  from trax.fastmath import numpy
  from trax.shapes import signature
label_desc: Use image node_id 12 for calling a global function with example usage: _custom_vjp(f, f_fwd, f_bwd, nondiff_argnums) and returns: _f, f