add primitive devices util

bayesflow/utils/__init__.py

@@ -5,6 +5,7 @@
 )
 from .callbacks import detailed_loss_callback
 from .comp_utils import expected_calibration_error
+from .devices import devices
 from .dict_utils import (
     convert_args,
     convert_kwargs,

bayesflow/utils/devices.py

@@ -0,0 +1,22 @@
+import keras
+
+
+def devices() -> list:
+    """Returns a list of available GPU devices."""
+    match keras.backend.backend():
+        case "jax":
+            import jax
+
+            return jax.devices("gpu")
+        case "tensorflow":
+            import tensorflow as tf
+
+            return tf.config.list_physical_devices("GPU")
+        case "torch":
+            import torch
+
+            return [torch.device(f"cuda:{i}") for i in range(torch.cuda.device_count())]
+        case "numpy":
+            return []
+        case _:
+            raise NotImplementedError(f"Backend {keras.backend.backend()} not supported.")