code: def primitive(f_raw):
      """
      Wraps a function so that its gradient can be specified and its invocation
      can be recorded. For examples, see the docs."""
  
      @wraps(f_raw)
      def f_wrapped(*args, **kwargs):
          boxed_args, trace, node_constructor = find_top_boxed_args(
              args
          )
          if boxed_args:
              argvals = subvals(
                  args,
                  [(argnum, box._value) for argnum, box in boxed_args],
              )
              if f_wrapped in notrace_primitives[node_constructor]:
                  return f_wrapped(*argvals, **kwargs)
              parents = tuple(box._node for _, box in boxed_args)
              argnums = tuple(argnum for argnum, _ in boxed_args)
              ans = f_wrapped(*argvals, **kwargs)
              node = node_constructor(
                  ans, f_wrapped, argvals, kwargs, argnums, parents
              )
              return new_box(ans, trace, node)
          else:
              return f_raw(*args, **kwargs)
  
      f_wrapped.fun = f_raw
      f_wrapped._is_autograd_primitive = True
      return f_wrapped
  
imports: import warnings
  from contextlib import contextmanager
  from collections import defaultdict
  from .util import subvals, toposort
  from .wrap_util import wraps
label_desc: Use image node_id 2 for calling a global function with example usage: primitive(f_raw) and returns: f_wrapped, new_box, f_raw, f_wrapped