code: def jacobian(fun, x):
      
      vjp, ans = _make_vjp(fun, x)
      ans_vspace = vspace(ans)
      jacobian_shape = ans_vspace.shape + vspace(x).shape
      grads = map(vjp, ans_vspace.standard_basis())
      return np.reshape(np.stack(grads), jacobian_shape)
docs: Returns a function which computes the Jacobian of `fun` with respect to
      positional argument number `argnum`, which must be a scalar or array. Unlike
      `grad` it is not restricted to scalar-output functions, but also it cannot
      take derivatives with respect to some argument types (like lists or dicts).
      If the input to `fun` has shape (in1, in2, ...) and the output has shape
      (out1, out2, ...) then the Jacobian has shape (out1, out2, ..., in1, in2, ...).