filter_vmap produces output with wrong out_axes arrangement

[remifan]

def foo(x):
    return x * 2.

x = jnp.arange(24).reshape((2,3,4))
y = jax.vmap(foo, out_axes=-1)(x)
z = eqx.filter_vmap(foo, out_axes=-1)(x)
print("jax:", y.shape, "eqx:", z.shape)
# jax: (3, 4, 2) eqx: (4, 3, 2)

equinox version: 0.11.8

[patrick-kidger]

Thank you for the report! I've just fixed this in #901.

This one is important enough that I'm also going to do a new Equinox release with the fix.