code: def elementwise_grad(fun, x):
      """
      Returns a function that computes the sum of each column of the Jacobian of
      `fun`, in one pass. If the Jacobian is diagonal, then this is the diagonal
      of the Jacobian.
      """
      vjp, ans = _make_vjp(fun, x)
      if vspace(ans).iscomplex:
          raise TypeError(
              "Elementwise_grad only applies to real-output functions."
          )
      return vjp(vspace(ans).ones())
  
imports: from __future__ import absolute_import
  from functools import partial
  from collections import OrderedDict
  import warnings
  from .wrap_util import unary_to_nary
  from .builtins import tuple
  from .core import make_vjp, make_jvp
  from .extend import primitive, defvjp_argnum, vspace
  import autograd.numpy
label_desc: Use image node_id 2 for calling a global function with example usage: elementwise_grad(fun, x) and returns: vjp