imt-ring 1.6.14__py3-none-any.whl → 1.6.15__py3-none-any.whl

{imt_ring-1.6.14.dist-info → imt_ring-1.6.15.dist-info}/METADATA

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

@@ -55,8 +55,8 @@ ring/ml/base.py,sha256=lfwEZLBDglOSRWChUHoH1kezefhttPV9TMEpNIqsMNw,9972
 ring/ml/callbacks.py,sha256=W19QF6_uvaNCjs8ObsjNXD7mv9gFgJBixdRSbB_BynE,13301
 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=pciltx7dR_yhFv2g8BvfV0VPgKm0HxXDxL2YjBPBaKQ,1400
+ring/ml/ringnet.py,sha256=Oud23uKmcvFtwNKdEu2KMMvNAFzJM_yBSRNz2a3CjL4,8670
+ring/ml/rnno_v1.py,sha256=2qE08OIvTJ5PvSxKpYGzGSrvEImWrdAT_qslZ7jP5tA,1372
 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.14.dist-info/METADATA,sha256=ZWWkJ_qO8O2KkW1m_HnXoz22H7ivmg7Qk4j7hEIgy9k,3821
-imt_ring-1.6.14.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-imt_ring-1.6.14.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.14.dist-info/RECORD,,
+imt_ring-1.6.15.dist-info/METADATA,sha256=zG-f_woph73I5ErczEeJaYXZLNahC-oNXDtzj26f1Po,3821
+imt_ring-1.6.15.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+imt_ring-1.6.15.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.15.dist-info/RECORD,,

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)

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)