imt-ring 1.5.0__tar.gz → 1.5.2__tar.gz

{imt_ring-1.5.0 → imt_ring-1.5.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.5.0
+Version: 1.5.2
 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.5.0 → imt_ring-1.5.2}/pyproject.toml

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

{imt_ring-1.5.0 → imt_ring-1.5.2}/src/imt_ring.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.5.0
+Version: 1.5.2
 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.5.0 → imt_ring-1.5.2}/src/ring/algorithms/_random.py

@@ -29,7 +29,8 @@ def random_angle_over_time(
     t_min: float,
     t_max: float | TimeDependentFloat,
     T: float,
-    Ts: float,
+    Ts: float | jax.Array,
+    N: Optional[int] = None,
     max_iter: int = 5,
     randomized_interpolation: bool = False,
     range_of_motion: bool = False,
@@ -84,7 +85,10 @@ def random_angle_over_time(
     )
 
     # resample
-    t = jnp.arange(T, step=Ts)
+    if N is None:
+        t = jnp.arange(T, step=Ts)
+    else:
+        t = jnp.arange(N) * Ts
     if randomized_interpolation:
         q = interpolate(cdf_bins_min, cdf_bins_max, method=interpolation_method)(
             t, ANG[:, 0], ANG[:, 1], consume
@@ -117,7 +121,8 @@ def random_position_over_time(
     t_max: float | TimeDependentFloat,
     T: float,
     Ts: float,
-    max_it: int,
+    N: Optional[int] = None,
+    max_it: int = 100,
     randomized_interpolation: bool = False,
     cdf_bins_min: int = 5,
     cdf_bins_max: Optional[int] = None,
@@ -203,7 +208,10 @@ def random_position_over_time(
     )
 
     # resample
-    t = jnp.arange(T, step=Ts)
+    if N is None:
+        t = jnp.arange(T, step=Ts)
+    else:
+        t = jnp.arange(N) * Ts
     if randomized_interpolation:
         r = interpolate(cdf_bins_min, cdf_bins_max, method=interpolation_method)(
             t, POS[:, 0], POS[:, 1], consume

{imt_ring-1.5.0 → imt_ring-1.5.2}/src/ring/algorithms/custom_joints/rr_imp_joint.py

@@ -2,6 +2,7 @@ from dataclasses import replace
 
 import jax
 import jax.numpy as jnp
+
 import ring
 from ring import maths
 from ring.algorithms.jcalc import _draw_rxyz
@@ -21,12 +22,12 @@ def register_rr_imp_joint(
         rot = ring.maths.quat_mul(rot_res, rot_pri)
         return ring.Transform.create(rot=rot)
 
-    def _draw_rr_imp(config, key_t, key_value, dt, _):
+    def _draw_rr_imp(config, key_t, key_value, dt, N, _):
         key_t1, key_t2 = jax.random.split(key_t)
         key_value1, key_value2 = jax.random.split(key_value)
-        q_traj_pri = _draw_rxyz(config, key_t1, key_value1, dt, _)
+        q_traj_pri = _draw_rxyz(config, key_t1, key_value1, dt, N, _)
         q_traj_res = _draw_rxyz(
-            replace(config_res, T=config.T), key_t2, key_value2, dt, _
+            replace(config_res, T=config.T), key_t2, key_value2, dt, N, _
         )
         # scale to be within bounds
         q_traj_res = q_traj_res * (jnp.deg2rad(ang_max_deg) / jnp.pi)

{imt_ring-1.5.0 → imt_ring-1.5.2}/src/ring/algorithms/custom_joints/suntay.py

@@ -225,7 +225,8 @@ def register_suntay(sconfig: SuntayConfig, name: str = "suntay"):
         mconfig: ring.MotionConfig,
         key_t: jax.random.PRNGKey,
         key_value: jax.random.PRNGKey,
-        dt: float,
+        dt: float | jax.Array,
+        N: int | None,
         _: jax.Array,
     ) -> jax.Array:
         key_value, consume = jax.random.split(key_value)
@@ -251,6 +252,7 @@ def register_suntay(sconfig: SuntayConfig, name: str = "suntay"):
             mconfig.t_max,
             mconfig.T,
             dt,
+            N,
             5,
             mconfig.randomized_interpolation_angle,
             mconfig.range_of_motion_hinge,

{imt_ring-1.5.0 → imt_ring-1.5.2}/src/ring/algorithms/generator/base.py

@@ -1,3 +1,4 @@
+import random
 from typing import Callable, Optional
 import warnings
 
@@ -33,8 +34,10 @@ class RCMG:
         randomize_positions: bool = False,
         randomize_motion_artifacts: bool = False,
         randomize_joint_params: bool = False,
+        randomize_hz: bool = False,
+        randomize_hz_kwargs: dict = dict(),
         imu_motion_artifacts: bool = False,
-        imu_motion_artifacts_kwargs: dict = dict(hide_injected_bodies=True),
+        imu_motion_artifacts_kwargs: dict = dict(),
         dynamic_simulation: bool = False,
         dynamic_simulation_kwargs: dict = dict(),
         output_transform: Optional[Callable] = None,
@@ -50,9 +53,6 @@ class RCMG:
         for c in config:
             assert c.is_feasible()
 
-        if cor:
-            sys = [s._replace_free_with_cor() for s in sys]
-
         self.gens = []
         for _sys in sys:
             self.gens.append(
@@ -71,6 +71,8 @@ class RCMG:
                     randomize_positions=randomize_positions,
                     randomize_motion_artifacts=randomize_motion_artifacts,
                     randomize_joint_params=randomize_joint_params,
+                    randomize_hz=randomize_hz,
+                    randomize_hz_kwargs=randomize_hz_kwargs,
                     imu_motion_artifacts=imu_motion_artifacts,
                     imu_motion_artifacts_kwargs=imu_motion_artifacts_kwargs,
                     dynamic_simulation=dynamic_simulation,
@@ -78,6 +80,7 @@ class RCMG:
                     output_transform=output_transform,
                     keep_output_extras=keep_output_extras,
                     use_link_number_in_Xy=use_link_number_in_Xy,
+                    cor=cor,
                 )
             )
 
@@ -174,35 +177,37 @@ class RCMG:
         sizes: int | list[int] = 1,
         seed: int = 1,
         shuffle: bool = True,
+        transform=None,
     ) -> types.BatchedGenerator:
         data = self.to_list(sizes, seed)
         assert len(data) >= batchsize
-
-        def data_fn(indices: list[int]):
-            return tree_utils.tree_batch([data[i] for i in indices])
-
-        return batch.generator_from_data_fn(
-            data_fn, list(range(len(data))), shuffle, batchsize
-        )
+        return self.eager_gen_from_list(data, batchsize, shuffle, transform)
 
     @staticmethod
-    def eager_gen_from_paths(
-        paths: str | list[str],
+    def eager_gen_from_list(
+        data: list[tree_utils.PyTree],
         batchsize: int,
-        include_samples: Optional[list[int]] = None,
         shuffle: bool = True,
-        load_all_into_memory: bool = False,
-        tree_transform=None,
-    ) -> tuple[types.BatchedGenerator, int]:
-        paths = utils.to_list(paths)
-        return batch.generator_from_paths(
-            paths,
-            batchsize,
-            include_samples,
-            shuffle,
-            load_all_into_memory=load_all_into_memory,
-            tree_transform=tree_transform,
-        )
+        transform=None,
+    ) -> types.BatchedGenerator:
+        data = data.copy()
+        n_batches, i = len(data) // batchsize, 0
+
+        def generator(key: jax.Array):
+            nonlocal i
+            if shuffle and i == 0:
+                random.shuffle(data)
+
+            start, stop = i * batchsize, (i + 1) * batchsize
+            batch = tree_utils.tree_batch(data[start:stop], backend="numpy")
+            batch = utils.pytree_deepcopy(batch)
+            if transform is not None:
+                batch = transform(batch)
+
+            i = (i + 1) % n_batches
+            return batch
+
+        return generator
 
 
 def _copy_dicts(f) -> dict:
@@ -231,6 +236,8 @@ def _build_mconfig_batched_generator(
     randomize_positions: bool,
     randomize_motion_artifacts: bool,
     randomize_joint_params: bool,
+    randomize_hz: bool,
+    randomize_hz_kwargs: dict,
     imu_motion_artifacts: bool,
     imu_motion_artifacts_kwargs: dict,
     dynamic_simulation: bool,
@@ -238,6 +245,7 @@ def _build_mconfig_batched_generator(
     output_transform: Callable | None,
     keep_output_extras: bool,
     use_link_number_in_Xy: bool,
+    cor: bool,
 ) -> types.BatchedGenerator:
 
     if add_X_jointaxes or add_y_relpose or add_y_rootincl:
@@ -284,13 +292,17 @@ def _build_mconfig_batched_generator(
         for f in pipe:
             key, consume = jax.random.split(key)
             sys = f(consume, sys)
+        if cor:
+            sys = sys._replace_free_with_cor()
         return sys
 
     def _finalize_fn(Xy: types.Xy, extras: types.OutputExtras):
         pipe = []
         if dynamic_simulation:
             pipe.append(finalize_fns.DynamicalSimulation(**dynamic_simulation_kwargs))
-        if imu_motion_artifacts and imu_motion_artifacts_kwargs["hide_injected_bodies"]:
+        if imu_motion_artifacts and imu_motion_artifacts_kwargs.get(
+            "hide_injected_bodies", True
+        ):
             pipe.append(motion_artifacts.HideInjectedBodies())
         if finalize_fn is not None:
             pipe.append(finalize_fns.FinalizeFn(finalize_fn))
@@ -312,19 +324,32 @@ def _build_mconfig_batched_generator(
         return Xy, extras
 
     def _gen(key: types.PRNGKey):
+        key, *consume = jax.random.split(key, len(config) + 1)
+        syss = jax.vmap(_setup_fn, (0, None))(jnp.array(consume), sys)
+
+        if randomize_hz:
+            assert "sampling_rates" in randomize_hz_kwargs
+            hzs = randomize_hz_kwargs["sampling_rates"]
+            assert len(set([c.T for c in config])) == 1
+            N = int(min(hzs) * config[0].T)
+            key, consume = jax.random.split(key)
+            dt = 1 / jax.random.choice(consume, jnp.array(hzs))
+            # makes sys.dt from float to AbstractArray
+            syss = syss.replace(dt=jnp.array(dt))
+        else:
+            N = None
+
         qs = []
-        for _config in config:
-            key, _q = draw_random_q(key, sys, _config)
+        for i, _config in enumerate(config):
+            key, _q = draw_random_q(key, syss[i], _config, N)
             qs.append(_q)
         qs = jnp.stack(qs)
 
-        key, *consume = jax.random.split(key, len(config) + 1)
-        syss = jax.vmap(_setup_fn, (0, None))(jnp.array(consume), sys)
-
         @jax.vmap
         def _vmapped_context(key, q, sys):
             x, _ = jax.vmap(kinematics.forward_kinematics_transforms, (None, 0))(sys, q)
-            Xy, extras = ({}, {}), (key, q, x, sys)
+            X = {"dt": jnp.array(sys.dt)} if randomize_hz else {}
+            Xy, extras = (X, {}), (key, q, x, sys)
             return _finalize_fn(Xy, extras)
 
         keys = jax.random.split(key, len(config))
@@ -340,6 +365,7 @@ def draw_random_q(
     key: types.PRNGKey,
     sys: base.System,
     config: jcalc.MotionConfig,
+    N: int | None,
 ) -> tuple[types.Xy, types.OutputExtras]:
 
     key_start = key
@@ -360,7 +386,7 @@ def draw_random_q(
         draw_fn = jcalc.get_joint_model(link_type).rcmg_draw_fn
         if draw_fn is None:
             raise Exception(f"The joint type {link_type} has no draw fn specified.")
-        q_link = draw_fn(config, key_t, key_value, sys.dt, joint_params)
+        q_link = draw_fn(config, key_t, key_value, sys.dt, N, joint_params)
         # even revolute and prismatic joints must be 2d arrays
         q_link = q_link if q_link.ndim == 2 else q_link[:, None]
         q_list.append(q_link)

imt_ring-1.5.2/src/ring/algorithms/generator/batch.py

@@ -0,0 +1,86 @@
+import jax
+import jax.numpy as jnp
+import numpy as np
+from tqdm import tqdm
+import tree_utils
+
+from ring import utils
+from ring.algorithms.generator import types
+
+
+def _build_batch_matrix(batchsizes: list[int]) -> jax.Array:
+    arr = []
+    for i, l in enumerate(batchsizes):
+        arr += [i] * l
+    return jnp.array(arr)
+
+
+def generators_lazy(
+    generators: list[types.BatchedGenerator],
+    repeats: list[int],
+    jit: bool = True,
+) -> types.BatchedGenerator:
+
+    batch_arr = _build_batch_matrix(repeats)
+    bs_total = len(batch_arr)
+    pmap, vmap = utils.distribute_batchsize(bs_total)
+    batch_arr = batch_arr.reshape((pmap, vmap))
+
+    pmap_trafo = jax.pmap
+    # single GPU node, then do jit + vmap instead of pmap
+    # this allows e.g. better NAN debugging capabilities
+    if pmap == 1:
+        pmap_trafo = lambda f: jax.jit(jax.vmap(f))
+    if not jit:
+        pmap_trafo = lambda f: jax.vmap(f)
+
+    @pmap_trafo
+    @jax.vmap
+    def _generator(key, which_gen: int):
+        return jax.lax.switch(which_gen, generators, key)
+
+    def generator(key):
+        pmap_vmap_keys = jax.random.split(key, bs_total).reshape((pmap, vmap, 2))
+        data = _generator(pmap_vmap_keys, batch_arr)
+
+        # merge pmap and vmap axis
+        data = utils.merge_batchsize(data, pmap, vmap, third_dim_also=True)
+        return data
+
+    return generator
+
+
+def generators_eager_to_list(
+    generators: list[types.BatchedGenerator],
+    n_calls: list[int],
+    seed: int = 1,
+    disable_tqdm: bool = False,
+) -> list[tree_utils.PyTree]:
+
+    key = jax.random.PRNGKey(seed)
+    data = []
+    for gen, n_call in tqdm(
+        zip(generators, n_calls),
+        desc="executing generators",
+        total=len(generators),
+        disable=disable_tqdm,
+    ):
+        for _ in tqdm(
+            range(n_call),
+            desc="number of calls for each generator",
+            total=n_call,
+            leave=False,
+            disable=disable_tqdm,
+        ):
+            key, consume = jax.random.split(key)
+            sample = gen(consume)
+            # converts also to numpy; but with np.array.flags.writeable = False
+            sample = jax.device_get(sample)
+            # this then sets this flag to True
+            sample = jax.tree_map(np.array, sample)
+
+            sample_flat, _ = jax.tree_util.tree_flatten(sample)
+            size = 1 if len(sample_flat) == 0 else sample_flat[0].shape[0]
+            data.extend([jax.tree_map(lambda a: a[i], sample) for i in range(size)])
+
+    return data

{imt_ring-1.5.0 → imt_ring-1.5.2}/src/ring/algorithms/generator/finalize_fns.py

@@ -251,8 +251,8 @@ def _expand_dt(X: dict, T: int):
     return X
 
 
-def _expand_then_flatten(args):
-    X, y = args
+def _expand_then_flatten(Xy):
+    X, y = Xy
     gyr = X["0"]["gyr"]
 
     batched = True

{imt_ring-1.5.0 → imt_ring-1.5.2}/src/ring/algorithms/generator/motion_artifacts.py

@@ -1,3 +1,4 @@
+import inspect
 import warnings
 
 import jax
@@ -127,6 +128,7 @@ def setup_fn_randomize_damping_stiffness_factory(
     prob_rigid: float = 0.0,
     all_imus_either_rigid_or_flex: bool = False,
     imus_surely_rigid: list[str] = [],
+    **kwargs,
 ):
     assert 0 <= prob_rigid <= 1
     assert prob_rigid != 1, "Use `imu_motion_artifacts`=False instead."
@@ -198,6 +200,18 @@ def setup_fn_randomize_damping_stiffness_factory(
     return setup_fn_randomize_damping_stiffness
 
 
+# assert that there exists no keyword arg duplicate which would induce ambiguity
+kwargs = lambda f: set(inspect.signature(f).parameters.keys())
+assert (
+    len(
+        kwargs(inject_subsystems).intersection(
+            kwargs(setup_fn_randomize_damping_stiffness_factory)
+        )
+    )
+    == 1
+)
+
+
 def _match_q_x_between_sys(
     sys_small: base.System,
     q_large: jax.Array,

{imt_ring-1.5.0 → imt_ring-1.5.2}/src/ring/algorithms/jcalc.py

@@ -274,8 +274,15 @@ def join_motionconfigs(
 
 
 DRAW_FN = Callable[
-    # config, key_t, key_value, dt, params
-    [MotionConfig, jax.random.PRNGKey, jax.random.PRNGKey, float, jax.Array],
+    # config, key_t, key_value, dt, N, params
+    [
+        MotionConfig,
+        jax.random.PRNGKey,
+        jax.random.PRNGKey,
+        float | jax.Array,
+        int | None,
+        jax.Array,
+    ],
     jax.Array,
 ]
 P_CONTROL_TERM = Callable[
@@ -410,7 +417,8 @@ def _draw_rxyz(
     config: MotionConfig,
     key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     _: jax.Array,
     # TODO, delete these args and pass a modifified `config` with `replace` instead
     enable_range_of_motion: bool = True,
@@ -435,6 +443,7 @@ def _draw_rxyz(
         config.t_max,
         config.T,
         dt,
+        N,
         max_iter,
         config.randomized_interpolation_angle,
         config.range_of_motion_hinge if enable_range_of_motion else False,
@@ -449,7 +458,8 @@ def _draw_pxyz(
     config: MotionConfig,
     _: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     __: jax.Array,
     cor: bool = False,
 ) -> jax.Array:
@@ -467,6 +477,7 @@ def _draw_pxyz(
         config.cor_t_max if cor else config.t_max,
         config.T,
         dt,
+        N,
         max_iter,
         config.randomized_interpolation_position,
         config.cdf_bins_min,
@@ -479,7 +490,8 @@ def _draw_spherical(
     config: MotionConfig,
     key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     _: jax.Array,
 ) -> jax.Array:
     # NOTE: We draw 3 euler angles and then build a quaternion.
@@ -491,6 +503,7 @@ def _draw_spherical(
             key_t,
             key_value,
             dt,
+            N,
             None,
             enable_range_of_motion=False,
             free_spherical=True,
@@ -506,7 +519,8 @@ def _draw_saddle(
     config: MotionConfig,
     key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     _: jax.Array,
 ) -> jax.Array:
     @jax.vmap
@@ -516,6 +530,7 @@ def _draw_saddle(
             key_t,
             key_value,
             dt,
+            N,
             None,
             enable_range_of_motion=False,
             free_spherical=False,
@@ -530,11 +545,12 @@ def _draw_p3d_and_cor(
     config: MotionConfig,
     _: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     __: jax.Array,
     cor: bool,
 ) -> jax.Array:
-    pos = jax.vmap(lambda key: _draw_pxyz(config, None, key, dt, None, cor))(
+    pos = jax.vmap(lambda key: _draw_pxyz(config, None, key, dt, N, None, cor))(
         jax.random.split(key_value, 3)
     )
     return pos.T
@@ -544,22 +560,24 @@ def _draw_p3d(
     config: MotionConfig,
     _: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     __: jax.Array,
 ) -> jax.Array:
-    return _draw_p3d_and_cor(config, _, key_value, dt, None, cor=False)
+    return _draw_p3d_and_cor(config, _, key_value, dt, N, None, cor=False)
 
 
 def _draw_cor(
     config: MotionConfig,
     _: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     __: jax.Array,
 ) -> jax.Array:
     key_value1, key_value2 = jax.random.split(key_value)
-    q_free = _draw_free(config, _, key_value1, dt, None)
-    q_p3d = _draw_p3d_and_cor(config, _, key_value2, dt, None, cor=True)
+    q_free = _draw_free(config, _, key_value1, dt, N, None)
+    q_p3d = _draw_p3d_and_cor(config, _, key_value2, dt, N, None, cor=True)
     return jnp.concatenate((q_free, q_p3d), axis=1)
 
 
@@ -567,12 +585,13 @@ def _draw_free(
     config: MotionConfig,
     key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     __: jax.Array,
 ) -> jax.Array:
     key_value1, key_value2 = jax.random.split(key_value)
-    q = _draw_spherical(config, key_t, key_value1, dt, None)
-    pos = _draw_p3d(config, None, key_value2, dt, None)
+    q = _draw_spherical(config, key_t, key_value1, dt, N, None)
+    pos = _draw_p3d(config, None, key_value2, dt, N, None)
     return jnp.concatenate((q, pos), axis=1)
 
 
@@ -580,7 +599,8 @@ def _draw_free_2d(
     config: MotionConfig,
     key_t: jax.random.PRNGKey,
     key_value: jax.random.PRNGKey,
-    dt: float,
+    dt: float | jax.Array,
+    N: int | None,
     __: jax.Array,
 ) -> jax.Array:
     key_value1, key_value2 = jax.random.split(key_value)
@@ -589,16 +609,20 @@ def _draw_free_2d(
         key_t,
         key_value1,
         dt,
+        N,
         None,
         enable_range_of_motion=False,
         free_spherical=True,
     )[:, None]
-    pos_yz = _draw_p3d(config, None, key_value2, dt, None)[:, :2]
+    pos_yz = _draw_p3d(config, None, key_value2, dt, N, None)[:, :2]
     return jnp.concatenate((angle_x, pos_yz), axis=1)
 
 
-def _draw_frozen(config: MotionConfig, _, __, dt: float, ___) -> jax.Array:
-    N = int(config.T / dt)
+def _draw_frozen(
+    config: MotionConfig, _, __, dt: float | jax.Array, N: int | None, ___
+) -> jax.Array:
+    if N is None:
+        N = int(config.T / dt)
     return jnp.zeros((N, 0))
 
 

{imt_ring-1.5.0 → imt_ring-1.5.2}/src/ring/ml/ml_utils.py

@@ -3,17 +3,16 @@ from functools import partial
 import os
 from pathlib import Path
 import pickle
-import random
 import time
 from typing import Optional, Protocol
 import warnings
 
 import jax
 import numpy as np
-import ring
-from ring.utils import import_lib
 from tree_utils import PyTree
 
+import ring
+from ring.utils import import_lib
 import wandb
 
 # An arbitrarily nested dictionary with Array leaves; Or strings
@@ -231,42 +230,5 @@ def save_model_tf(jax_func, path: str, *input, validate: bool = True):
         )
 
 
-def train_val_split(
-    tps: list[str],
-    bs: int,
-    n_batches_for_val: int = 1,
-    transform_gen=None,
-    tree_transform=None,
-):
-    "Uses `random` module for shuffeling."
-    if transform_gen is None:
-        transform_gen = lambda gen: gen
-
-    len_val = n_batches_for_val * bs
-
-    _, N = ring.RCMG.eager_gen_from_paths(tps, 1)
-    include_samples = list(range(N))
-    random.shuffle(include_samples)
-
-    train_data, val_data = include_samples[:-len_val], include_samples[-len_val:]
-    X_val, y_val = transform_gen(
-        ring.RCMG.eager_gen_from_paths(
-            tps, len_val, val_data, tree_transform=tree_transform
-        )[0]
-    )(jax.random.PRNGKey(420))
-
-    generator = transform_gen(
-        ring.RCMG.eager_gen_from_paths(
-            tps,
-            bs,
-            train_data,
-            load_all_into_memory=True,
-            tree_transform=tree_transform,
-        )[0]
-    )
-
-    return generator, (X_val, y_val)
-
-
 def _unknown_link_names(N: int):
     return [f"link{i}" for i in range(N)]

{imt_ring-1.5.0 → imt_ring-1.5.2}/src/ring/utils/__init__.py

@@ -16,6 +16,7 @@ from .utils import pickle_load
 from .utils import pickle_save
 from .utils import primes
 from .utils import pytree_deepcopy
+from .utils import replace_elements_w_nans
 from .utils import sys_compare
 from .utils import to_list
 from .utils import tree_equal