undefined f if vmap applied to lu.WrappedFun
In `vmap`, `f` is not be defined whenever the if condition is false:

https://github.com/google/jax/blob/f8eda9bb645a1ed489f08be57bae380bbcf5fa0c/jax/api.py#L414-L420
diff --git a/jax/api.py b/jax/api.py
index 03f925880..ed4652ed6 100644
--- a/jax/api.py
+++ b/jax/api.py
@@ -413,8 +413,7 @@ def vmap(fun, in_axes=0, out_axes=0):
 
   @wraps(fun, docstr=docstr)
   def batched_fun(*args, **kwargs):
-    if not isinstance(fun, lu.WrappedFun):
-      f = lu.wrap_init(fun, kwargs)
+    f = lu.wrap_init(fun, kwargs) if not isinstance(fun, lu.WrappedFun) else fun
     in_axes_ = in_axes if isinstance(in_axes, (list, tuple)) else (in_axes,) * len(args)
     in_flat, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
     jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(f, in_trees)
['jax/api.py']