Repository: google/jax

Issue:
'TypeError: iteration over a 0-d array' when using vmap
When I'm using vmap as shown in the following example I get a 'TypeError: iteration over a 0-d array' but expect an array with the second derivatives for both input variables. I'am using jax 0.2.10 and jaxlib 0.1.64.

```
import jax.numpy as jnp
from jax import grad, vmap


def target_function(x, y):
    single_input = jnp.array([x, y])
    return jnp.sum(single_input ** 3)

single_input = jnp.array([1., 1.])

vmap_second_derivatives = (
    lambda variable_position: grad(grad(target_function, variable_position), variable_position)(*single_input)
)

variable_positions = jnp.array([0, 1])

second_derivatives = vmap(vmap_second_derivatives)(variable_positions)
```

  
**The full error message:**

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Traceback (most recent call last):
  File "/Users/davidanton/develop/issue_vmap/issueVMAP.py", line 18, in <module>
    second_derivatives = vmap(vmap_second_derivatives)(variable_positions)
  File "/Users/davidanton/develop/issue_vmap/issueVMAP.py", line 13, in <lambda>
    lambda variable_position: grad(grad(target_function, variable_position), variable_position)(*single_input)
jax._src.traceback_util.FilteredStackTrace: TypeError: iteration over a 0-d array

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/davidanton/develop/issue_vmap/issueVMAP.py", line 18, in <module>
    second_derivatives = vmap(vmap_second_derivatives)(variable_positions)
  File "/Users/davidanton/develop/issue_vmap/venv/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/davidanton/develop/issue_vmap/venv/lib/python3.9/site-packages/jax/api.py", line 1237, in batched_fun
    out_flat = batching.batch(
  File "/Users/davidanton/develop/issue_vmap/venv/lib/python3.9/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/davidanton/develop/issue_vmap/issueVMAP.py", line 13, in <lambda>
    lambda variable_position: grad(grad(target_function, variable_position), variable_position)(*single_input)
  File "/Users/davidanton/develop/issue_vmap/venv/lib/python3.9/site-packages/jax/api.py", line 748, in grad
    value_and_grad_f = value_and_grad(fun, argnums, has_aux=has_aux,
  File "/Users/davidanton/develop/issue_vmap/venv/lib/python3.9/site-packages/jax/api.py", line 808, in value_and_grad
    argnums = _ensure_index(argnums)
  File "/Users/davidanton/develop/issue_vmap/venv/lib/python3.9/site-packages/jax/api_util.py", line 38, in _ensure_index
    return tuple(map(operator.index, x))
  File "/Users/davidanton/develop/issue_vmap/venv/lib/python3.9/site-packages/jax/_src/util.py", line 37, in safe_map
    args = list(map(list, args))
  File "/Users/davidanton/develop/issue_vmap/venv/lib/python3.9/site-packages/jax/core.py", line 497, in __iter__
    return iter(self.aval._iter(self))
  File "/Users/davidanton/develop/issue_vmap/venv/lib/python3.9/site-packages/jax/_src/lax/lax.py", line 1943, in _iter
    raise TypeError("iteration over a 0-d array")  # same as numpy error
TypeError: iteration over a 0-d array


**Calling the function without vmap, on the other hand, works fine.**

```
import jax.numpy as jnp
from jax import grad, vmap


def target_function(x, y):
    single_input = jnp.array([x, y])
    return jnp.sum(single_input ** 3)

single_input = jnp.array([1., 1.])

vmap_second_derivatives = (
    lambda variable_position: grad(grad(target_function, variable_position), variable_position)(*single_input)
)

print(vmap_second_derivatives(0))
print(vmap_second_derivatives(1))
```

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
6.0
6.0


Discussion:

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -768,7 +768,7 @@ def value_and_grad(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
             "same shape as the arguments at positions {argnums}.")
 
   _check_callable(fun)
-  argnums = _ensure_index(argnums)
+  argnums = core.concrete_or_error(_ensure_index, argnums)
 
   @wraps(fun, docstr=docstr, argnums=argnums)
   @api_boundary


Test patch: