imt-ring 1.6.12__py3-none-any.whl → 1.6.13__py3-none-any.whl

{imt_ring-1.6.12.dist-info → imt_ring-1.6.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.12
+Version: 1.6.13
 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.6.12.dist-info → imt_ring-1.6.13.dist-info}/RECORD

@@ -53,10 +53,10 @@ ring/io/xml/to_xml.py,sha256=fohb-jWMf2cxVdT5dmknsGyrNMseICSbKEz_urbaWbQ,3407
 ring/ml/__init__.py,sha256=nbh48gaswWeY4S4vT1sply_3ROj2DQ7agjoLR4Ho3T8,1517
 ring/ml/base.py,sha256=lfwEZLBDglOSRWChUHoH1kezefhttPV9TMEpNIqsMNw,9972
 ring/ml/callbacks.py,sha256=W19QF6_uvaNCjs8ObsjNXD7mv9gFgJBixdRSbB_BynE,13301
-ring/ml/ml_utils.py,sha256=siiRWbUpjYQz1nAlARm47oqR2K74YTiE1syCoOEmiWw,6370
+ring/ml/ml_utils.py,sha256=1GXJfeoXbwCbRdYA2np3CbJpSupaw4eyf3quh9y4BO0,6462
 ring/ml/optimizer.py,sha256=fWyF__ezUltrA16SLfOC1jvS3zBh9NJsMYa6-V0frhs,4709
 ring/ml/ringnet.py,sha256=Tb2WJ_cc5L3mk1lo0NOfkpXIzJZXf4PJ5aLPtHQyUmY,8650
-ring/ml/rnno_v1.py,sha256=T4SKG7iypqn2HBQLKhDmJ2Slj2Z5jtUBHvX_6aL8pyM,1103
+ring/ml/rnno_v1.py,sha256=ujyIkDxMSTag9iRFEmoHqfqSrlOFjcZs9_rBbLd8p9Q,1380
 ring/ml/train.py,sha256=huUfMK6eotS6BRrQKoZ-AUG0um3jlqpfQFZNJT8LKiE,10854
 ring/ml/training_loop.py,sha256=CEokvPQuuk_WCd-J60ZDodJYcPVvyxLfgXDr_DnbzRI,3359
 ring/ml/params/0x13e3518065c21cd8.pickle,sha256=Zh2k1zK-TNxJl5F7nyTeQ9001qqRE_dfvaq1HWV287A,9355838
@@ -84,7 +84,7 @@ ring/utils/randomize_sys.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3
 ring/utils/utils.py,sha256=tJaWXLGOTwkxJQj2l23dX97wO3aZYhM2qd7eNuMRs84,6907
 ring/utils/register_gym_envs/__init__.py,sha256=PtPIRBQJ16339xZ9G9VpvqrvcGbQ_Pk_SUz4tQPa9nQ,94
 ring/utils/register_gym_envs/saddle.py,sha256=tA5CyW_akSXyDm0xJ83CtOrUMVElH0f9vZtEDDJQalI,4422
-imt_ring-1.6.12.dist-info/METADATA,sha256=NIcGCBCzA9jwqxvyHYHl5QdfiaFLLxdnQjOk17YX0bA,3821
-imt_ring-1.6.12.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-imt_ring-1.6.12.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.12.dist-info/RECORD,,
+imt_ring-1.6.13.dist-info/METADATA,sha256=wMwfHX8PsYaxXZRldgqd71fGltIfUo_W9xVE9DSz5o0,3821
+imt_ring-1.6.13.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+imt_ring-1.6.13.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.13.dist-info/RECORD,,

ring/ml/ml_utils.py

@@ -184,6 +184,8 @@ def on_cluster() -> bool:
 
 
 def unique_id() -> str:
+    if wandb.run is not None:
+        wandb.config.setdefault("unique_id", ring._UNIQUE_ID)
     return ring._UNIQUE_ID
 
 

ring/ml/rnno_v1.py

@@ -4,6 +4,8 @@ import haiku as hk
 import jax
 import jax.numpy as jnp
 
+from .ringnet import LSTM
+
 
 def rnno_v1_forward_factory(
     output_dim: int,
@@ -13,18 +15,29 @@ def rnno_v1_forward_factory(
     act_fn_linear=jax.nn.relu,
     act_fn_rnn=jax.nn.elu,
     lam: Optional[tuple[int]] = None,
+    celltype: str = "gru",
 ):
     # unused
     del lam
 
+    if celltype == "gru":
+        _cell = hk.GRU
+        _factor = 1
+    elif celltype == "lstm":
+        _cell = LSTM
+        _factor = 2
+    else:
+        raise NotImplementedError
+
     @hk.without_apply_rng
     @hk.transform_with_state
     def forward_fn(X):
         assert X.shape[-2] == 1
 
         for i, n_units in enumerate(rnn_layers):
+            n_units = _factor * n_units
             state = hk.get_state(f"rnn_{i}", shape=[1, n_units], init=jnp.zeros)
-            X, state = hk.dynamic_unroll(hk.GRU(n_units), X, state)
+            X, state = hk.dynamic_unroll(_cell(n_units), X, state)
             hk.set_state(f"rnn_{i}", state)
 
             if layernorm: