From 921bca5105c180a9d2971efe7e1f8fb08c6f83ef Mon Sep 17 00:00:00 2001
From: Google DeepMind <noreply@deepmind.com>
Date: Mon, 21 Jul 2025 09:46:49 -0700
Subject: [PATCH] #mjx don't specify JAX default backend

PiperOrigin-RevId: 785472058
Change-Id: Iba2d3e87dd2d78e6eee8fb579a760250a91b5d23
---
 mjx/mujoco/mjx/_src/io.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index f5a66995..1eada20e 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -148,7 +148,7 @@ def _resolve_impl_and_device(
   elif (has_impl, has_device) == (False, True):
     impl = _resolve_impl(device)
   else:
-    device = jax.devices(jax.default_backend())[0]
+    device = jax.devices()[0]
     logging.info('Using JAX default device: %s.', device)
     impl = _resolve_impl(device)