# Buggy code
def minimize(
    fun: Callable,
    x0: jnp.ndarray,
    args: Tuple = (),
    *,
    method: str,
    tol: Optional[float] = None,
    options: Optional[Mapping[str, Any]] = None,
) -> OptimizeResults:
    """Minimization of scalar function of one or more variables.

    This API for this function matches SciPy with some minor deviations:

    - Gradients of ``fun`` are calculated automatically using JAX's autodiff
      support when required.
    - The ``method`` argument is required. You must specify a solver.
    - Various optional arguments in the SciPy interface have not yet been
      implemented.
    - Optimization results may differ from SciPy due to differences in the line
      search implementation.

    ``minimize`` supports ``jit`` compilation. It does not yet support
    differentiation or arguments in the form of multi-dimensional arrays, but
    support for both is planned.

    Args:
      fun: the objective function to be minimized, ``fun(x, *args) -> float``,
        where ``x`` is an 1-D array with shape ``(n,)`` and ``args`` is a tuple
        of the fixed parameters needed to completely specify the function.
        ``fun`` must support differentiation.
      x0: initial guess. Array of real elements of size ``(n,)``, where 'n' is
        the number of independent variables.
      args: extra arguments passed to the objective function.
      method: solver type. Currently only "BFGS" is supported.
      tol: tolerance for termination. For detailed control, use solver-specific
        options.
      options: a dictionary of solver options. All methods accept the following
        generic options:
          maxiter : int
              Maximum number of iterations to perform. Depending on the
              method each iteration may use several function evaluations.

    Returns: OptimizeResults object.
    """
    if options is None:
        options = {}

    fun_with_args = lambda x: fun(x, *args)

    if method.lower() == "bfgs":
        results = minimize_bfgs(fun_with_args, x0, **options)
        message = (
            "status meaning: 0=converged, 1=max BFGS iters reached, "
            "3=zoom failed, 4=saddle point reached, "
            "5=max line search iters reached, -1=undefined"
        )
        success = (results.converged) & jnp.logical_not(results.failed)
        return OptimizeResults(
            x=results.x_k,
            success=success,
            status=results.status,
            message=message,
            fun=results.f_k,
            jac=results.g_k,
            hess_inv=results.H_k,
            nfev=results.nfev,
            njev=results.ngev,
            nit=results.k,
        )

    raise ValueError("Method {} not recognized".format(method))


# Traceback (ValueError)
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, [*,*,*,*,*,*,*,*,*,*,*]).

# Fixed code
def minimize(
    fun: Callable,
    x0: jnp.ndarray,
    args: Tuple = (),
    *,
    method: str,
    tol: Optional[float] = None,
    options: Optional[Mapping[str, Any]] = None,
) -> OptimizeResults:
    """Minimization of scalar function of one or more variables.

    This API for this function matches SciPy with some minor deviations:

    - Gradients of ``fun`` are calculated automatically using JAX's autodiff
      support when required.
    - The ``method`` argument is required. You must specify a solver.
    - Various optional arguments in the SciPy interface have not yet been
      implemented.
    - Optimization results may differ from SciPy due to differences in the line
      search implementation.

    ``minimize`` supports ``jit`` compilation. It does not yet support
    differentiation or arguments in the form of multi-dimensional arrays, but
    support for both is planned.

    Args:
      fun: the objective function to be minimized, ``fun(x, *args) -> float``,
        where ``x`` is an 1-D array with shape ``(n,)`` and ``args`` is a tuple
        of the fixed parameters needed to completely specify the function.
        ``fun`` must support differentiation.
      x0: initial guess. Array of real elements of size ``(n,)``, where 'n' is
        the number of independent variables.
      args: extra arguments passed to the objective function.
      method: solver type. Currently only "BFGS" is supported.
      tol: tolerance for termination. For detailed control, use solver-specific
        options.
      options: a dictionary of solver options. All methods accept the following
        generic options:
          maxiter : int
              Maximum number of iterations to perform. Depending on the
              method each iteration may use several function evaluations.

    Returns: OptimizeResults object.
    """
    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":
        results = minimize_bfgs(fun_with_args, x0, **options)
        message = (
            "status meaning: 0=converged, 1=max BFGS iters reached, "
            "3=zoom failed, 4=saddle point reached, "
            "5=max line search iters reached, -1=undefined"
        )
        success = (results.converged) & jnp.logical_not(results.failed)
        return OptimizeResults(
            x=results.x_k,
            success=success,
            status=results.status,
            message=message,
            fun=results.f_k,
            jac=results.g_k,
            hess_inv=results.H_k,
            nfev=results.nfev,
            njev=results.ngev,
            nit=results.k,
        )

    raise ValueError("Method {} not recognized".format(method))