Skip to content

Enable mjw on MacOS #2948

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hartikainen wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Enable mjw on MacOS #2948

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b48275f
Enable `mjw` on `MacOS`
hartikainen Nov 16, 2025
8962f3a
Enable jax-warp ffi for cpu
hartikainen Feb 1, 2026
62b486c
Fix jax-warp ffi
hartikainen Mar 19, 2026
8e8d151
Fix unstable test constraint sorting on CPU warp
hartikainen Mar 19, 2026
877f79a
Enable io tests on non-cuda device
hartikainen Mar 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 14 additions & 14 deletions mjx/mujoco/mjx/_src/io.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,23 +90,20 @@ def _resolve_device(
logging.debug('Picking default device: %s.', device_0)
return device_0

if impl == types.Impl.WARP:
if has_cuda_gpu_device():
cuda_gpus = jax.devices('cuda')
logging.debug('Picking default device: %s', cuda_gpus[0])
device_0 = cuda_gpus[0]
else:
device_0 = jax.devices('cpu')[0]
return device_0

if impl == types.Impl.CPP:
cpu_0 = jax.devices('cpu')[0]
logging.debug('Picking default device: %s', cpu_0)
return cpu_0

if impl == types.Impl.WARP:
# WARP implementation requires a CUDA GPU.
cuda_gpus = [d for d in jax.devices('cuda')]
if not cuda_gpus:
raise AssertionError(
'No CUDA GPU devices found in'
f' jax.devices("cuda")={jax.devices("cuda")}.'
)

logging.debug('Picking default device: %s', cuda_gpus[0])
return cuda_gpus[0]

raise ValueError(f'Unsupported implementation: {impl}')


Expand All @@ -121,9 +118,12 @@ def _check_impl_device_compatibility(
impl = types.Impl(impl)

if impl == types.Impl.WARP:
if not _is_cuda_gpu_device(device):
is_cuda_device = _is_cuda_gpu_device(device)
is_cpu_device = device.platform == 'cpu'
if not (is_cuda_device or is_cpu_device):
raise AssertionError(
f'Warp implementation requires a CUDA GPU device, got {device}.'
'Warp implementation requires a CUDA GPU or CPU device, got '
f'{device}.'
)
_check_warp_installed()

Expand Down
37 changes: 5 additions & 32 deletions mjx/mujoco/mjx/_src/io_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,8 +142,6 @@ def setUp(self):
def test_put_model(self, xml, impl):
if impl == 'warp' and not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if impl == 'warp' and not mjx_io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')

m = mujoco.MjModel.from_xml_string(xml)
mx = mjx.put_model(m, impl=impl)
Expand Down Expand Up @@ -311,8 +309,6 @@ def test_put_model_warp_has_expected_shapes(self):
"""Tests that put_model produces expected shapes for MuJoCo Warp."""
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not mjx_io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')

m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
mx = mjx.put_model(m, impl='warp')
Expand All @@ -335,8 +331,6 @@ def test_put_model_warp_graph_mode(self, mode: str | None):
"""Tests that put_model accepts graph_mode parameter."""
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp not installed.')
if not mjx_io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device available.')

if mode is None:
graph_mode = None
Expand Down Expand Up @@ -472,8 +466,6 @@ def test_make_data(self, impl: str):
def test_make_data_warp(self):
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp is not installed.')
if not mjx_io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device.')
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
d = mjx.make_data(m, impl='warp', nconmax=9, njmax=23)
self.assertEqual(d._impl.contact__dist.shape[0], 9)
Expand All @@ -485,8 +477,6 @@ def test_put_data(self, impl: str):
if impl == 'warp':
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp is not installed.')
if not mjx_io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device.')

m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
d = mujoco.MjData(m)
Expand Down Expand Up @@ -604,8 +594,6 @@ def test_put_data_warp_ndim(self):
"""Tests that put_data produces expected dimensions for Warp fields."""
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp is not installed.')
if not mjx_io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device.')

m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
d = mujoco.MjData(m)
Expand Down Expand Up @@ -770,8 +758,6 @@ def test_get_data_into_warp(self):
# and remove this test.
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp is not installed.')
if not mjx_io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device.')

m = mujoco.MjModel.from_xml_string('<mujoco></mujoco>')
d = mujoco.MjData(m)
Expand Down Expand Up @@ -847,8 +833,6 @@ def test_make_data_warp_has_expected_shapes(self):
"""Tests that make_data produces expected shapes for MuJoCo Warp."""
if not mjxw.WARP_INSTALLED:
self.skipTest('Warp is not installed.')
if not mjx_io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device.')

m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
dx = mjx.make_data(m, impl='warp')
Expand All @@ -871,8 +855,6 @@ def test_data_slice(self, impl):
"""Tests that slice on Data works as expected."""
if impl == 'warp' and not mjxw.WARP_INSTALLED:
self.skipTest('Warp is not installed.')
if impl == 'warp' and not mjx_io.has_cuda_gpu_device():
self.skipTest('No CUDA GPU device.')

m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
dx = jax.vmap(lambda x: mjx.make_data(m, impl=impl))(jp.arange(10))
Expand Down Expand Up @@ -934,8 +916,8 @@ def put_data(dummy_arg_for_batching):
('gpu-nvidia', 'jax', ('gpu', Impl.JAX)),
('tpu', 'jax', ('tpu', Impl.JAX)),
# WARP backend specified.
('cpu', 'warp', ('cpu', 'error')),
('gpu-notnvidia', 'warp', ('cpu', 'error')),
('cpu', 'warp', ('cpu', Impl.WARP)),
('gpu-notnvidia', 'warp', ('gpu', 'error')),
('gpu-nvidia', 'warp', ('gpu', Impl.WARP)),
('tpu', 'warp', ('tpu', 'error')),
# CPP backend specified.
Expand All @@ -962,10 +944,10 @@ def put_data(dummy_arg_for_batching):
('gpu-nvidia', 'jax', ('gpu', Impl.JAX)),
('tpu', 'jax', ('tpu', Impl.JAX)),
# WARP backend impl specified.
('cpu', 'warp', ('cpu', 'error')),
('gpu-notnvidia', 'warp', ('cpu', 'error')),
('cpu', 'warp', ('cpu', Impl.WARP)),
('gpu-notnvidia', 'warp', ('cpu', Impl.WARP)),
('gpu-nvidia', 'warp', ('gpu', Impl.WARP)),
('tpu', 'warp', ('tpu', 'error')),
('tpu', 'warp', ('cpu', Impl.WARP)),
# CPP backend impl specified, CPU should always be available.
('cpu', 'cpp', ('cpu', Impl.CPP)),
('gpu-notnvidia', 'cpp', ('cpu', Impl.CPP)),
Expand Down Expand Up @@ -1140,15 +1122,6 @@ def backends_side_effect():
self.mock_jax_backends.side_effect = backends_side_effect

expected_device, expected_impl = expected
if (
expected_impl == 'error'
and default_device_str != 'gpu-nvidia'
and impl_str == 'warp'
):
with self.assertRaisesRegex(RuntimeError, 'cuda backend not supported'):
mjx_io._resolve_impl_and_device(impl=impl_str, device=None)
return

if expected_impl == 'error':
with self.assertRaises(AssertionError):
mjx_io._resolve_impl_and_device(impl=impl_str, device=None)
Expand Down
Loading
Loading