Worth a shot if someone wants to code it up. A couple of thoughts:

• This is probably going to matter most for nonlinear models and will likely require passing additional information to the simulators (eg, gradients) to take advantage of the structure, when available. In addition, we may have to modify the interconnect function and perhaps a few others to use JAX instead of NumPy for various operations.
• If we implement this, we should probably use a x_method keyword (with 'x' as some appropriate prefix) so that users can choose to use 'jax' (if installed) versus 'numpy'. This would also let us maintain use of NumPy if JAX is not installed (so that JAX is an optional dependency, like pandas and cvxopt).
• It would be good to put together some benchmarks so we can make sure things are faster. See the benchmarks/ directory for examples and the README file in that directory for more information.
• Another possibility (and method) would be to allow the use of numba. This can possibly be done completely by the user, though again if you want to interconnect subsystems then some modifications to nlsys.py may be required.