imt-ring 1.3.9__py3-none-any.whl → 1.3.10__py3-none-any.whl

{imt_ring-1.3.9.dist-info → imt_ring-1.3.10.dist-info}/METADATA

@@ -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.dist-info → imt_ring-1.3.10.dist-info}/RECORD

@@ -15,7 +15,7 @@ ring/algorithms/custom_joints/rr_joint.py,sha256=jnRtjtOCALMaq2_0bcu2d7qgfQ6etXp
 ring/algorithms/custom_joints/suntay.py,sha256=7-kym1kMDwqYD_2um1roGcBeB8BlTCPe1wljuNGNARA,16676
 ring/algorithms/generator/__init__.py,sha256=p4ucl0zQtp5NwNoXIRjmTzGGRu2WOAWFfNmYRPwQles,912
 ring/algorithms/generator/base.py,sha256=sr-YZkjd8pZJAI5vFG_IqOO4AEeiEYtXr8uUsPMS6Q4,14779
-ring/algorithms/generator/batch.py,sha256=kNlq78W-nAtbp6Xe82UjbPY-rXX2alGLxTokTITSbAc,9226
+ring/algorithms/generator/batch.py,sha256=bslFSN2Gs_aX9cNwFooExhKUwevc70q3bspEMTwygm4,9256
 ring/algorithms/generator/motion_artifacts.py,sha256=_kiAl1VHoX1fW5AUlXOtPBWyHIIFof_M78AP-m9f1ME,8790
 ring/algorithms/generator/pd_control.py,sha256=XJ_Gd5AkIRh-jBrMfQyMXjVwhx2gCNHznjzFbmAwhZs,5767
 ring/algorithms/generator/randomize.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3699
@@ -71,14 +71,15 @@ ring/sys_composer/__init__.py,sha256=5J_JJJIHfTPcpxh0v4FqiOs81V1REPUd7pgiw2nAN5E
 ring/sys_composer/delete_sys.py,sha256=cIM9KbyLfg7B9121g7yjzuFbjeNu9cil1dPavAYEgzk,3408
 ring/sys_composer/inject_sys.py,sha256=Mj-q-mUjXKwkg-ol6IQAjf9IJfk7pGhez0_WoTKTgm0,3503
 ring/sys_composer/morph_sys.py,sha256=2GpPtS5hT0eZMptdGpt30Hc97OykJNE67lEVRf7sHrc,12700
-ring/utils/__init__.py,sha256=rTvSA4RiJAVCY_A64FUMd8IJTv94LgoSA3Ps5X63_jA,799
-ring/utils/batchsize.py,sha256=mPFGD7AedFMycHtyIuZtNWCaAvKLLWSWaB7X6u54xvM,1358
+ring/utils/__init__.py,sha256=FZ9ziQrWlx16QIpQ8RdLKrvN_17CAdvnZMNNodxWY0o,812
+ring/utils/backend.py,sha256=cKSi9sB59texqKzNVASTDczGKLCBL8VVDiP7TNdj41k,1294
+ring/utils/batchsize.py,sha256=3sLrmQMvCyoJ0Qi1uVzApSMfleYrIFAG7Y5_KuhJ-k4,1554
 ring/utils/colab.py,sha256=ZLHwP0jNQUsmZJU4l68a5djULPi6T-jYNNHevjIoMn8,1631
 ring/utils/hdf5.py,sha256=BzXwVypZmEZeHVgeGZ78YYdi10NEQtnPhdrb8dQAXo0,5856
 ring/utils/normalizer.py,sha256=67L2BU1MRsMT4pD41ta3JJMppLN0ozFmnwrmXDtnqrQ,1698
 ring/utils/path.py,sha256=hAfSlqRi-ew536RnjDDM7IKapdMJc-EvhrR0Y-BCFWc,1265
 ring/utils/utils.py,sha256=mIcKNv5v2de8HrG7bAhl2bNfmwkMZyIIwFkJq2XWMOI,5357
-imt_ring-1.3.9.dist-info/METADATA,sha256=H65-QICwM4mtRPumYJbrenN74nmiMBGbeV3pecKEeOg,3104
-imt_ring-1.3.9.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-imt_ring-1.3.9.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.3.9.dist-info/RECORD,,
+imt_ring-1.3.10.dist-info/METADATA,sha256=4z1vTo1XeXzJVPcAHPa9Y8psVMFF4BB2lpm9G82aADY,3105
+imt_ring-1.3.10.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+imt_ring-1.3.10.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.3.10.dist-info/RECORD,,

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

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

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
+    )

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")