<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 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())