Use jax.sharding instead of deprecated backend

ethanluoyc · Feb 22, 2024 · 197acb7 · 197acb7
1 parent fa68a2f
commit 197acb7
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/corax/agents/jax/actors.py b/corax/agents/jax/actors.py
@@ -65,8 +65,14 @@ def __init__(
 
         # Unpack ActorCore, jitting if requested.
         if jit:
-            self._init = jax.jit(actor.init, backend=backend)
-            self._policy = jax.jit(actor.select_action, backend=backend)
+            if backend is not None:
+                sharding = jax.sharding.SingleDeviceSharding(
+                    jax.local_devices(backend=backend)[0]
+                )
+            else:
+                sharding = None
+            self._init = jax.jit(actor.init, sharding)  # type: ignore
+            self._policy = jax.jit(actor.select_action, sharding)  # type: ignore
         else:
             self._init = actor.init
             self._policy = actor.select_action