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

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

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

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: imt-ring
-Version: 1.6.14
+Version: 1.6.15
 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.14 → imt_ring-1.6.15}/src/ring/ml/ringnet.py

@@ -175,7 +175,7 @@ class LSTM(hk.RNNCore):
         prev_state: jax.Array,
     ):
         if len(inputs.shape) > 2 or not inputs.shape:
-            raise ValueError("LSTM input must be rank-1 or rank-2.")
+            raise ValueError(f"LSTM input must be rank-1 or rank-2; not {inputs.shape}")
         prev_state_h = prev_state[: self.hidden_size]
         prev_state_c = prev_state[self.hidden_size :]
         x_and_h = jnp.concatenate([inputs, prev_state_h], axis=-1)

{imt_ring-1.6.14 → imt_ring-1.6.15}/src/ring/ml/rnno_v1.py

@@ -33,13 +33,12 @@ def rnno_v1_forward_factory(
     @hk.transform_with_state
     def forward_fn(X):
         assert X.shape[-2] == 1
+        X = X[..., 0, :]
 
         for i, n_units in enumerate(rnn_layers):
-            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])
+            state = hk.get_state(f"rnn_{i}", shape=[n_units * _factor], init=jnp.zeros)
+            X, state = hk.dynamic_unroll(_cell(n_units), X, state)
+            hk.set_state(f"rnn_{i}", state)
 
             if layernorm:
                 X = hk.LayerNorm(axis=-1, create_scale=False, create_offset=False)(X)