imt-ring 1.3.9__tar.gz → 1.3.10__tar.gz

{imt_ring-1.3.9 → imt_ring-1.3.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.3.9
+Version: 1.3.10
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.3.9 → imt_ring-1.3.10}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "imt-ring"
-version = "1.3.9"
+version = "1.3.10"
 authors = [
   { name="Simon Bachhuber", email="simon.bachhuber@fau.de" },
 ]

{imt_ring-1.3.9 → imt_ring-1.3.10}/src/imt_ring.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.3.9
+Version: 1.3.10
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.3.9 → imt_ring-1.3.10}/src/imt_ring.egg-info/SOURCES.txt

@@ -80,6 +80,7 @@ src/ring/sys_composer/delete_sys.py
 src/ring/sys_composer/inject_sys.py
 src/ring/sys_composer/morph_sys.py
 src/ring/utils/__init__.py
+src/ring/utils/backend.py
 src/ring/utils/batchsize.py
 src/ring/utils/colab.py
 src/ring/utils/hdf5.py

{imt_ring-1.3.9 → imt_ring-1.3.10}/src/ring/algorithms/generator/batch.py

@@ -63,12 +63,12 @@ def batch_generators_lazy(
 
 
 def _number_of_executions_required(size: int) -> int:
-    vmap_threshold = 128
     _, vmap = utils.distribute_batchsize(size)
 
+    eager_threshold = utils.batchsize_thresholds()[1]
     primes = iter(utils.primes(vmap))
     n_calls = 1
-    while vmap > vmap_threshold:
+    while vmap > eager_threshold:
         prime = next(primes)
         n_calls *= prime
         vmap /= prime

{imt_ring-1.3.9 → imt_ring-1.3.10}/src/ring/utils/__init__.py

@@ -1,4 +1,4 @@
-from .batchsize import backend
+from .batchsize import batchsize_thresholds
 from .batchsize import distribute_batchsize
 from .batchsize import expand_batchsize
 from .batchsize import merge_batchsize

imt_ring-1.3.10/src/ring/utils/backend.py

@@ -0,0 +1,30 @@
+import os
+import re
+
+
+def set_host_device_count(n):
+    """
+    By default, XLA considers all CPU cores as one device. This utility tells XLA
+    that there are `n` host (CPU) devices available to use. As a consequence, this
+    allows parallel mapping in JAX :func:`jax.pmap` to work in CPU platform.
+
+    .. note:: This utility only takes effect at the beginning of your program.
+        Under the hood, this sets the environment variable
+        `XLA_FLAGS=--xla_force_host_platform_device_count=[num_devices]`, where
+        `[num_device]` is the desired number of CPU devices `n`.
+
+    .. warning:: Our understanding of the side effects of using the
+        `xla_force_host_platform_device_count` flag in XLA is incomplete. If you
+        observe some strange phenomenon when using this utility, please let us
+        know through our issue or forum page. More information is available in this
+        `JAX issue <https://github.com/google/jax/issues/1408>`_.
+
+    :param int n: number of CPU devices to use.
+    """
+    xla_flags = os.getenv("XLA_FLAGS", "")
+    xla_flags = re.sub(
+        r"--xla_force_host_platform_device_count=\S+", "", xla_flags
+    ).split()
+    os.environ["XLA_FLAGS"] = " ".join(
+        ["--xla_force_host_platform_device_count={}".format(n)] + xla_flags
+    )

{imt_ring-1.3.9 → imt_ring-1.3.10}/src/ring/utils/batchsize.py

@@ -1,19 +1,37 @@
-from typing import Optional, Tuple
+from typing import Tuple, TypeVar
 
 import jax
-from tree_utils import PyTree
+
+PyTree = TypeVar("PyTree")
+
+
+def batchsize_thresholds():
+    backend = jax.default_backend()
+    if backend == "cpu":
+        vmap_size_min = 1
+        eager_threshold = 4
+    elif backend == "gpu":
+        vmap_size_min = 8
+        eager_threshold = 128
+    else:
+        raise Exception(
+            f"Backend {backend} has no default values, please add them in this function"
+        )
+    return vmap_size_min, eager_threshold
 
 
 def distribute_batchsize(batchsize: int) -> Tuple[int, int]:
     """Distributes batchsize accross pmap and vmap."""
-    vmap_size_min = 8
+    vmap_size_min = batchsize_thresholds()[0]
     if batchsize <= vmap_size_min:
         return 1, batchsize
     else:
         n_devices = jax.local_device_count()
-        assert (
-            batchsize % n_devices
-        ) == 0, f"Your GPU count of {n_devices} does not split batchsize {batchsize}"
+        msg = (
+            f"Your local device count of {n_devices} does not split batchsize"
+            + f" {batchsize}. local devices are {jax.local_devices()}"
+        )
+        assert (batchsize % n_devices) == 0, msg
         vmap_size = int(batchsize / n_devices)
         return int(batchsize / vmap_size), vmap_size
 
@@ -35,17 +53,3 @@ def expand_batchsize(tree: PyTree, pmap_size: int, vmap_size: int) -> PyTree:
         ),
         tree,
     )
-
-
-CPU_ONLY = False
-
-
-def backend(cpu_only: bool = False, n_gpus: Optional[int] = None):
-    "Sets backend for all jax operations (including this library)."
-    global CPU_ONLY
-
-    if cpu_only and not CPU_ONLY:
-        CPU_ONLY = True
-        from jax import config
-
-        config.update("jax_platform_name", "cpu")