Get JAX device from string

Toni-SM · Jun 22, 2024 · 1ad585e · 1ad585e
1 parent 4a9bb1e
commit 1ad585e
Show file tree

Hide file tree

Showing 7 changed files with 47 additions and 16 deletions.
diff --git a/skrl/agents/jax/base.py b/skrl/agents/jax/base.py
@@ -53,7 +53,10 @@ def __init__(self,
         if device is None:
             self.device = jax.devices()[0]
         else:
-            self.device = device if isinstance(device, jax.Device) else jax.devices(device)[0]
+            self.device = device
+            if type(device) == str:
+                device_type, device_index = f"{device}:0".split(':')[:2]
+                self.device = jax.devices(device_type)[int(device_index)]
 
         if type(memory) is list:
             self.memory = memory[0]

diff --git a/skrl/envs/wrappers/jax/base.py b/skrl/envs/wrappers/jax/base.py
@@ -20,12 +20,19 @@ def __init__(self, env: Any) -> None:
         self._env = env
 
         # device (faster than @property)
-        self.device = jax.devices()[0]
+        self.device = None
         if hasattr(self._env, "device"):
-            try:
-                self.device = jax.devices(self._env.device.split(':')[0] if type(self._env.device) == str else self._env.device.type)[0]
-            except RuntimeError:
-                pass
+            if type(self._env.device) == str:
+                device_type, device_index = f"{self._env.device}:0".split(':')[:2]
+                try:
+                    self.device = jax.devices(device_type)[int(device_index)]
+                except RuntimeError:
+                    self.device = None
+            else:
+                self.device = self._env.device
+        if self.device is None:
+            self.device = jax.devices()[0]
+
         # spaces
         try:
             self._action_space = self._env.single_action_space
@@ -135,12 +142,18 @@ def __init__(self, env: Any) -> None:
         self._env = env
 
         # device (faster than @property)
-        self.device = jax.devices()[0]
+        self.device = None
         if hasattr(self._env, "device"):
-            try:
-                self.device = jax.devices(self._env.device.split(':')[0] if type(self._env.device) == str else self._env.device.type)[0]
-            except RuntimeError:
-                pass
+            if type(self._env.device) == str:
+                device_type, device_index = f"{self._env.device}:0".split(':')[:2]
+                try:
+                    self.device = jax.devices(device_type)[int(device_index)]
+                except RuntimeError:
+                    self.device = None
+            else:
+                self.device = self._env.device
+        if self.device is None:
+            self.device = jax.devices()[0]
 
         self.possible_agents = []
 

diff --git a/skrl/memories/jax/base.py b/skrl/memories/jax/base.py
@@ -70,7 +70,10 @@ def __init__(self,
         if device is None:
             self.device = jax.devices()[0]
         else:
-            self.device = device if isinstance(device, jax.Device) else jax.devices(device)[0]
+            self.device = device
+            if type(device) == str:
+                device_type, device_index = f"{device}:0".split(':')[:2]
+                self.device = jax.devices(device_type)[int(device_index)]
 
         # internal variables
         self.filled = False

diff --git a/skrl/models/jax/base.py b/skrl/models/jax/base.py
@@ -79,7 +79,10 @@ def __call__(self, inputs, role):
         if device is None:
             self.device = jax.devices()[0]
         else:
-            self.device = device if isinstance(device, jax.Device) else jax.devices(device)[0]
+            self.device = device
+            if type(device) == str:
+                device_type, device_index = f"{device}:0".split(':')[:2]
+                self.device = jax.devices(device_type)[int(device_index)]
 
         self.observation_space = observation_space
         self.action_space = action_space

diff --git a/skrl/multi_agents/jax/base.py b/skrl/multi_agents/jax/base.py
@@ -60,7 +60,10 @@ def __init__(self,
         if device is None:
             self.device = jax.devices()[0]
         else:
-            self.device = device if isinstance(device, jax.Device) else jax.devices(device)[0]
+            self.device = device
+            if type(device) == str:
+                device_type, device_index = f"{device}:0".split(':')[:2]
+                self.device = jax.devices(device_type)[int(device_index)]
 
         # convert the models to their respective device
         for _models in self.models.values():

diff --git a/skrl/resources/noises/jax/base.py b/skrl/resources/noises/jax/base.py
@@ -31,7 +31,10 @@ def sample(self, size):
         if device is None:
             self.device = jax.devices()[0]
         else:
-            self.device = device if isinstance(device, jax.Device) else jax.devices(device)[0]
+            self.device = device
+            if type(device) == str:
+                device_type, device_index = f"{device}:0".split(':')[:2]
+                self.device = jax.devices(device_type)[int(device_index)]
 
     def sample_like(self, tensor: Union[np.ndarray, jax.Array]) -> Union[np.ndarray, jax.Array]:
         """Sample a noise with the same size (shape) as the input tensor

diff --git a/skrl/resources/preprocessors/jax/running_standard_scaler.py b/skrl/resources/preprocessors/jax/running_standard_scaler.py
@@ -95,7 +95,10 @@ def __init__(self,
         if device is None:
             self.device = jax.devices()[0]
         else:
-            self.device = device if isinstance(device, jax.Device) else jax.devices(device)[0]
+            self.device = device
+            if type(device) == str:
+                device_type, device_index = f"{device}:0".split(':')[:2]
+                self.device = jax.devices(device_type)[int(device_index)]
 
         size = self._get_space_size(size)