Can't vmap across input using Gauss Newton fwd

Vmapping across y0 with any method using AbstractGaussNewton throws a TypeError. MWE

import jax
import jax.numpy as jnp

import optimistix as optx


def rosenbrock(x, args):
    del args
    term1 = 10 * (x[1:] - x[:-1] ** 2)
    term2 = x - 1
    return term1, term2


inits = jnp.zeros((4, 10))
solve = lambda x: optx.least_squares(rosenbrock, optx.LevenbergMarquardt(1e-8, 1e-9), x)
out = jax.vmap(solve)(inits)  # throws error

The reason of this looks to be that the state includes an f_info with a FunctionLinearOperator whose linearised function is a Jaxpr which can't be batched over.

Start with the provided JAX MWE and trace the AbstractGaussNewton state, especially f_info and its FunctionLinearOperator. Check how the linearised function is batched under jax.vmap; done means the example runs without a TypeError across y0.