imt-ring 1.6.12__tar.gz → 1.6.14__tar.gz

{imt_ring-1.6.12 → imt_ring-1.6.14}/PKG-INFO

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

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

{imt_ring-1.6.12 → imt_ring-1.6.14}/src/imt_ring.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.12
+Version: 1.6.14
 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 → imt_ring-1.6.14}/src/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
 
 

{imt_ring-1.6.12 → imt_ring-1.6.14}/src/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,19 +15,31 @@ 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):
-            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)
-            hk.set_state(f"rnn_{i}", state)
+            state = hk.get_state(
+                f"rnn_{i}", shape=[1, n_units * _factor], init=jnp.zeros
+            )
+            X, state = hk.dynamic_unroll(_cell(n_units), X[..., 0, :], state[0])
+            hk.set_state(f"rnn_{i}", state[None])
 
             if layernorm:
                 X = hk.LayerNorm(axis=-1, create_scale=False, create_offset=False)(X)