code: def _tf_grad(f, **kwargs):
      """Grad with support for argnums."""
      argnums = kwargs.pop("argnums", 0)
      if argnums != 0:
  
          def g(*args, **kwargs):
              args = list(args)
              args[0], args[argnums] = args[argnums], args[0]
              return f(*args, **kwargs)
  
      else:
          g = f
      grad_g = tf_np_extensions.grad(g, **kwargs)
      if argnums == 0:
          return grad_g
  
      def grad_f(*args, **kwargs):
          args = list(args)
          args[0], args[argnums] = args[argnums], args[0]
          return grad_g(*args, **kwargs)
  
      return grad_f
  
file_path: trax/trax/fastmath/tf.py
imports: import numpy
  import tensorflow.compat.v2
  from trax.shapes import ShapeDtype
  from trax.tf_numpy import extensions
  from trax.tf_numpy import numpy
label_desc: Use image node_id 7 for calling a global function with example usage: _tf_grad(f) and returns: grad_f, grad_g, grad_g, f