Google DeepMind
521d152ec5
This change modifies the Warp JAX FFI to register callbacks for both CUDA and CPU platforms, allowing Warp kernels to be launched on CPU devices. MJX's io.py is updated to allow Warp to use CPU devices if no CUDA GPU is available. Tests are adjusted to no longer skip when CUDA GPUs are absent. A minor fix for sorting precision in test_util.py is also included. Kristian Hartikainen - https://github.com/google-deepmind/mujoco/pull/2948 Github issue - https://github.com/google-deepmind/mujoco/issues/2947 PiperOrigin-RevId: 906018844 Change-Id: I9c026b4a4e2a0276d5d3af3954fd027cf75b6d7a