RK45 Jax Solver Returning NaN Solution

[keatincf]

The RK45 Jax solver doesn't always return a valid result. The current processing will produce a matrix of solutions from the solver and then utilize the last solution. However, sometimes the matrix of solutions aren't all valid solutions and a row in the solution matrix can just be all NaN. This can end up causing issues for clients and prevent stepping dynamics.

Another underlying issue is that the Jax solver is utilized just by having Jax in your Python environment. If a client isn't expecting to use Jax, but happens to have a dependency that pulls in Jax, they can have processing change without them knowing. Since the solver is still considered experimental, it probably makes sense to require the client to be deliberate about their choice to use that solver.

[keatincf]

I think this issue points to:

• Handling of error conditions arising from use of the Jax solver that aligns with the solver's API contract
• Prevent automated switching between Jax and non-jax processing, but require deliberate activation of Jax processing