code: def primitive(f_raw):
      
      @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
docs: Wraps a function so that its gradient can be specified and its invocation
      can be recorded. For examples, see the docs.