Remove asarray workaround?

[bagibence]

I just noticed that the JAX issue optimistix._misc.asarray was working around seems to be fixed now.
jax-ml/jax#18020 is still open, but it might not be a problem because jax>=0.7.2 "represents constants in its internal jaxpr representation as a TypedNdArray".

I added a non-array leaf to the smoke_aux used in the tests and the example from @patrick-kidger's comment. They pass with jax>=0.7.2, but not with earlier versions.

This example from another comment doesn't pass, but not sure how essential that is for Optimistix:

f, consts = jax.closure_convert(jax.numpy.asarray, 1.0)
out = f(1.0, *consts)
assert type(out) is not float

JAX commits that might be relevant:
jax-ml/jax@b43662f
jax-ml/jax@00538d0
jax-ml/jax@cb4fcd8

I don't have a good understanding of JAX internals, so please feel free to ignore this if it's not helpful.

[bagibence]

The CI fails because jax>=0.7.0 requires Python 3.11, so dropping support for 3.10 is another change here.

[johannahaffner]

Hi Bence!

Thank you for looking into this :) On your specific points:

• Dropping compatibility with earlier versions of JAX and correspondingly earlier versions of Equinox is not a price I'm willing to pay to drop a small utility function. We definitely have users whose code still uses older versions, and there are cases where older versions are more performant for certain specialised tasks.
• Dropping support for Python 3.10 is reasonable though - I think it makes sense to match what JAX is doing in that regard, plus we'd adhere to the "last three versions of Python" rule. Our code would continue to run on older versions to the extent it does today, assuming a compatible JAX version. We'd just no longer make the promise that we'll put in the maintenance work to make sure that it continues to do so. Could you edit this PR to do only this?

[bagibence]

Of course, but I'll rather open a tiny PR for dropping 3.10. Do you want one for updating smoke_aux or just drop that?

Ignoring the need to support older JAX versions, would this be a valid solution? I wanted to replicate the setup done in minimise without private functions from Optimistix or copying the workaround. Do you see anything that might go wrong here?

[johannahaffner]

Ignoring the need to support older JAX versions, would this be a valid solution? I wanted to replicate the setup done in minimise without private functions from Optimistix or copying the workaround. Do you see anything that might go wrong here?

No, this should work. On older versions the JVP of jnp.asarray was the problem, not the function itself. If you're no longer running into the same bug then this is probably a good way to go :)

[johannahaffner]

I'll close this PR :)