Repository: google/jax
Issue: Strange behavior for jax.scipy.optimize.minimize when using vmap

Hi all, thanks for developing such a fantastic package. I've been excited to use jax in my day to day research, as well as trainees' research in my group.

I'm experimenting with applying `vmap` to a scalar function and passing it along to `jax.scipy.optimize.minimize`. I'm seeing a strange error regarding `in_axes` specification for `vmap` after the fact. Here is code to reproduce the error below:

```python
import jax
import jax.numpy as jnp
import jax.scipy.optimize as sopt
import numpy as np

# set up keys
seed = 1234
rng_key = jax.random.PRNGKey(seed)

# split
rng_key, x_key, b_key, e_key = jax.random.split(rng_key, 4)

# gen random data and a binary outcome
X = jax.random.normal(x_key, shape=(10, 3))
B = jax.random.normal(b_key, shape=(3,2))
eps = jax.random.normal(e_key, shape=(10, 2))
tholds = jnp.array([0, 1])
Y = ((X @ B + eps) > tholds).astype(float)

# loglike
def loglike(b_hat, y):
    t1 = jnp.sum(-jnp.log1p(jnp.exp(X @ b_hat)), axis=0)
    t2 = jnp.sum(y * (X @ b_hat), axis=0)
    return t1 + t2

# nll for min
def nll(b_hat, y):
    return -loglike(b_hat, y)

# vmap to span dimensions
NLL = jax.vmap(nll, (1, 1), 0)

# test it out
print("testing nll = ", NLL(B, Y))

# run it
sopt.minimize(NLL, jnp.zeros((3, 2)), Y, method="BFGS")
```

```
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
testing nll =  [5.0628967 5.4723964]
Traceback (most recent call last):
  File "test.py", line 37, in <module>
    sopt.minimize(NLL, jnp.zeros((3, 2)), Y, method="BFGS")
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/scipy/optimize/minimize.py", line 96, in minimize
    results = minimize_bfgs(fun_with_args, x0, **options)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/scipy/optimize/bfgs.py", line 99, in minimize_bfgs
    f_0, g_0 = jax.value_and_grad(fun)(x0)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/scipy/optimize/minimize.py", line 93, in <lambda>
    fun_with_args = lambda x: fun(x, *args)
jax._src.traceback_util.FilteredStackTrace: ValueError: vmap in_axes specification must be a tree prefix of the corresponding value, got specification (1, 1) for value tree PyTreeDef(tuple, [*,*,*,*,*,*,*,*,*,*,*]).

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 "test.py", line 37, in <module>
    sopt.minimize(NLL, jnp.zeros((3, 2)), Y, method="BFGS")
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/scipy/optimize/minimize.py", line 96, in minimize
    results = minimize_bfgs(fun_with_args, x0, **options)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/scipy/optimize/bfgs.py", line 99, in minimize_bfgs
    f_0, g_0 = jax.value_and_grad(fun)(x0)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/api.py", line 808, in value_and_grad_f
    ans, vjp_py = _vjp(f_partial, *dyn_args)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/api.py", line 1914, in _vjp
    out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py", line 114, in vjp
    out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py", line 101, in linearize
    jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 506, in trace_to_jaxpr
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/scipy/optimize/minimize.py", line 93, in <lambda>
    fun_with_args = lambda x: fun(x, *args)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/api.py", line 1220, in batched_fun
    in_axes_flat = flatten_axes("vmap in_axes", in_tree, (in_axes, 0), kws=True)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/api_util.py", line 188, in flatten_axes
    raise ValueError(f"{name} specification must be a tree prefix of the "
ValueError: vmap in_axes specification must be a tree prefix of the corresponding value, got specification (1, 1) for value tree PyTreeDef(tuple, [*,*,*,*,*,*,*,*,*,*,*]).
```

There is some similarly strange behavior when trying to `vmap` over `minimize` directly,
```python
# try vmap across minimize directly
def my_min(b0, y):
    return sopt.minimize(nll, b0, y, method="BFGS")

vmin = jax.vmap(my_min, (1, 1), 0)
vmin(B, Y)
```

With output,
```
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
testing nll =  [5.0628967 5.4723964]
Traceback (most recent call last):
  File "test.py", line 44, in <module>
    vmin(B, Y)
  File "test.py", line 41, in my_min
    return sopt.minimize(nll, b0, y, method="BFGS")
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/scipy/optimize/minimize.py", line 96, in minimize
    results = minimize_bfgs(fun_with_args, x0, **options)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/scipy/optimize/bfgs.py", line 99, in minimize_bfgs
    f_0, g_0 = jax.value_and_grad(fun)(x0)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/scipy/optimize/minimize.py", line 93, in <lambda>
    fun_with_args = lambda x: fun(x, *args)
jax._src.traceback_util.FilteredStackTrace: TypeError: nll() takes 2 positional arguments but 11 were given

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 "test.py", line 44, in <module>
    vmin(B, Y)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/api.py", line 1222, in batched_fun
    out_flat = batching.batch(
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "test.py", line 41, in my_min
    return sopt.minimize(nll, b0, y, method="BFGS")
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/scipy/optimize/minimize.py", line 96, in minimize
    results = minimize_bfgs(fun_with_args, x0, **options)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/scipy/optimize/bfgs.py", line 99, in minimize_bfgs
    f_0, g_0 = jax.value_and_grad(fun)(x0)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/api.py", line 808, in value_and_grad_f
    ans, vjp_py = _vjp(f_partial, *dyn_args)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/api.py", line 1914, in _vjp
    out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py", line 114, in vjp
    out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py", line 101, in linearize
    jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 506, in trace_to_jaxpr
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/nicholas/opt/miniconda3/lib/python3.8/site-packages/jax/_src/scipy/optimize/minimize.py", line 93, in <lambda>
    fun_with_args = lambda x: fun(x, *args)
TypeError: nll() takes 2 positional arguments but 11 were given
```

Changed files:
tests/scipy_optimize_test.py, jax/_src/scipy/optimize/minimize.py

diff --git a/jax/_src/scipy/optimize/minimize.py b/jax/_src/scipy/optimize/minimize.py
index a21f61b27..eff99746f 100644
--- a/jax/_src/scipy/optimize/minimize.py
+++ b/jax/_src/scipy/optimize/minimize.py
@@ -91,6 +91,10 @@ def minimize(
   if options is None:
     options = {}
 
+  if not isinstance(args, tuple):
+    msg = "args argument to jax.scipy.optimize.minimize must be a tuple, got {}"
+    raise TypeError(msg.format(args))
+
   fun_with_args = lambda x: fun(x, *args)
 
   if method.lower() == 'bfgs':
diff --git a/tests/scipy_optimize_test.py b/tests/scipy_optimize_test.py
index 1c504f224..19a46e116 100644
--- a/tests/scipy_optimize_test.py
+++ b/tests/scipy_optimize_test.py
@@ -94,6 +94,13 @@ class TestBFGS(jtu.JaxTestCase):
     results = jax.scipy.optimize.minimize(f, jnp.ones(n), method='BFGS')
     self.assertAllClose(results.x, jnp.zeros(n), atol=1e-6, rtol=1e-6)
 
+  def test_args_must_be_tuple(self):
+    A = jnp.eye(2) * 1e4
+    def f(x):
+      return jnp.mean((A @ x) ** 2)
+    with self.assertRaisesRegex(TypeError, "args .* must be a tuple"):
+      jax.scipy.optimize.minimize(f, jnp.ones(2), args=45, method='BFGS')
+
 
 if __name__ == "__main__":
   absltest.main()