Adjust code for jax==0.7.1, set new version number

Author: markschoene

README.md

@@ -2,6 +2,12 @@
 ![Figure 1](docs/figure1.png)
 This is the official implementation of our paper [Scalable Event-by-event Processing of Neuromorphic Sensory Signals With Deep State-Space Models
 ](https://arxiv.org/abs/2404.18508).
+## Version 0.2 released
+- Fixed a bug in the event tokenization for DVS. Improves our DVS Gestures result to 99.2% accuracy (without using spatial convolutions at all)
+- Updated DVS configs
+- Compatibility with recent JAX versions >= 0.7.1
+
+## Introduction
 The core motivation for this work was the irregular time-series modeling problem presented in the paper [Simplified State Space Layers for Sequence Modeling
 ](https://arxiv.org/abs/2208.04933). 
 We acknowledge the awesome [S5 project](https://github.com/lindermanlab/S5) and the trainer class provided by this [UvA tutorial](https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/guide4/Research_Projects_with_JAX.html), which highly influenced our code.

event_ssm/train_utils.py

@@ -152,7 +152,7 @@ def get_first_device(x):
 def print_model_size(params, name=''):
     fn_is_complex = lambda x: x.dtype in [np.complex64, np.complex128]
     param_sizes = map_nested_fn(lambda k, param: param.size * (2 if fn_is_complex(param) else 1))(params)
-    total_params_size = sum(jax.tree_leaves(param_sizes))
+    total_params_size = sum(jax.tree.leaves(param_sizes))
     print('[*] Model parameter count:', total_params_size)
 
 

event_ssm/trainer.py

@@ -61,7 +61,7 @@ def __init__(
         self.best_eval_metrics = {}
 
         # logger details
-        self.log_dir = os.path.join(self.log_config.log_dir)
+        self.log_dir = os.path.abspath(os.path.join(self.log_config.log_dir))
         print('[*] Logging to', self.log_dir)
 
         if not os.path.isdir(self.log_dir):
@@ -72,7 +72,7 @@ def __init__(
             os.makedirs(os.path.join(self.log_dir, 'checkpoints'))
 
         num_parameters = int(sum(
-            [arr.size for arr in jax.tree_flatten(self.train_state.params)[0]
+            [arr.size for arr in jax.tree.flatten(self.train_state.params)[0]
              if isinstance(arr, Array)]
         ) / self.world_size)
         print("[*] Number of model parameters:", num_parameters)

setup.py

@@ -3,7 +3,7 @@
 setup(
     name='Event-based-SSM',
     packages=['event_ssm'],
-    version='0.1',
+    version='0.2',
     description='Event-stream modeling with state-space models',
     author='Mark Schoene',
     author_email='[email protected]',