NaNs encountered when taking the quotient of two arrays

Hi, I'm quite new to Optimistix (and JAX more broadly), so please forgive the simple question if the solution turns out to be obvious.

I'm attempting to solve a root finding problem which relies quite heavily on external function calls in the residual. These obviously require special treatment, but doesn't seem to actually be the cause of my problems in this instance. The key to this it seems is that taking the quotient of two arrays is producing NaN values in the JVP which propagate through the rest of my code.

import jax

jax.config.update("jax_disable_jit", True)
jax.config.update("jax_debug_nans", True)

class Solver(eqx.Module):
    arr: Float[Array, "n"]

    def residuals(self, y, args):
        z = self._quotient(y)
        return jnp.array([0.5, 0.5]) - z

    def solve(self):
        y0 = jnp.array([0.9, 0.1])
        solver = optx.LevenbergMarquardt(rtol=1e-5, atol=1e-5)

        sol = optx.root_find(
            self.residuals,
            solver,
            y0=y0,
            args=(),
        )

        return sol

    def _quotient(self, y):
        # FloatingPointError: invalid value (nan) encountered in div
        return y / self.arr


if __name__ == "__main__":
    problem = Solver(jnp.array([5.0, 2.5]))
    sol = problem.solve()
    print(sol.value)

For this toy example I can re-enable the JIT or disable NaN checking and everything works just fine. However, it seems that the NaN values persist, only that JAX is able to handle them appropriately. The problem still remains for me however, as these NaN values in the tangent will propagate through the pure callbacks and land in my external functions which do not know to handle them appropriately.

I am not really sure if this is a JAX or Optimistix (or more likely my) problem, but given that I can get a Jacobian out of the residual just fine I thought it would be best to start here.

Thanks for any help!

Start by running the supplied Python reproducer with jax_disable_jit and jax_debug_nans enabled, then trace the JVP through Solver._quotient and optx.root_find. Compare behavior with JIT and NaN checking disabled; done means isolating whether the tangent NaNs originate in Optimistix or JAX and documenting a reproducible fix or workaround for the external callbacks.