Compatibility with jaxtyping

[nstarman]

Hi! Really enjoying quax. I've been working to get galax potentials quaxified (most relevantly in GalacticDynamics/galax#187) and ran into a compatibility issue with jaxtyping's runtime type checking.

If a module has runtime type checking turned on, e.g. install_import_hook("galax.potential", "beartype.beartype") then quaxed functions don't pass objects through correctly. As an example from GalacticDynamics/galax#187

>>> from jax_quantity import Quantity
>>> import galax.potential as gp
>>> import galax.units as gu
>>> pot = gp.KeplerPotential(m=Quantity(1e12, "Msun"), units=gu.galactic)
>>> pot._potential_energy(Quantity([1.0, 0, 0], "kpc"), Quantity(0, "Myr"), pot._G)
TypeCheckError: Type-check error whilst checking the parameters of _potential_energy.
The problem arose whilst typechecking argument 'q'.
Called with arguments: {'self': KeplerPotential(...), 'q': f64[3], 't': i64[], '_G': f64[]}
Parameter annotations: (self, q: Shaped[Quantity, '*batch 3'], t: Union[Shaped[Quantity, '*#batch '], Shaped[Quantity, '*#batch ']], /, _G: Float[Quantity, '']).

[patrick-kidger]

This looks expected to me. Once a Value passes through a quaxify boundary, then it should look like an array to everything inside.

Put another way, the underlying function should be one that acts on arrays. It is by wrapping it in a quaxify that it becomes able to consume Values.

If you use the jaxtyping import hook, then the type-checking is always put on the bottom of the decorator list (this is needed for compatibility with jax.custom_{jvp,vjp}. In particular, that means it happens inside the quaxify.