[BUG] Warp modules used in JAX FFI graph captures can be unloaded prematurely

### Bug Description

We should retain the modules with kernels that JAX captured in a CUDA graph. This problem manifested as module unloading during graph capture here:
https://github.com/google-deepmind/mujoco_warp/pull/350


### System Information

_No response_