code: def elementwise_grad(fun, x):
      
      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())
docs: 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.