brainstate 0.1.0.post20241220__py2.py3-none-any.whl → 0.1.0.post20241221__py2.py3-none-any.whl

brainstate/nn/_dynamics/_state_delay.py

@@ -51,11 +51,14 @@ def _get_delay(delay_time, delay_step):
             assert isinstance(delay_step, int), '"delay_step" should be an integer.'
             if delay_step == 0:
                 return 0., 0
-            delay_time = delay_step * environ.get_dt()
+            with jax.ensure_compile_time_eval():
+                delay_time = delay_step * environ.get_dt()
     else:
         assert delay_step is None, '"delay_step" should be None if "delay_time" is given.'
         # assert isinstance(delay_time, (int, float))
-        delay_step = math.ceil(delay_time / environ.get_dt())
+        with jax.ensure_compile_time_eval():
+            delay_step = delay_time / environ.get_dt()
+        delay_step = math.ceil(float(delay_step))
     return delay_time, delay_step
 
 

{brainstate-0.1.0.post20241220.dist-info → brainstate-0.1.0.post20241221.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20241220
+Version: 0.1.0.post20241221
 Summary: A ``State``-based Transformation System for Program Compilation and Augmentation.
 Home-page: https://github.com/chaobrain/brainstate
 Author: BrainState Developers

{brainstate-0.1.0.post20241220.dist-info → brainstate-0.1.0.post20241221.dist-info}/RECORD

@@ -93,7 +93,7 @@ brainstate/nn/_dynamics/__init__.py,sha256=j1HSWu01wf5-KjSaNhBC9utVGDALOhUsFPrLP
 brainstate/nn/_dynamics/_dynamics_base.py,sha256=SOfrET1Z-JkS2GPo2bkQ-v40FLozIxT5A7ObwW_hQYo,21987
 brainstate/nn/_dynamics/_dynamics_base_test.py,sha256=gXMwENqqSvyZbMpLP0QtYndJ_h39dF5gIeiiSbMAjTk,2721
 brainstate/nn/_dynamics/_projection_base.py,sha256=jYe3WdBMgz2TJkcxPWEkyK7OA4IR1ChISd2GTfM6U2o,13528
-brainstate/nn/_dynamics/_state_delay.py,sha256=nZYGmVKmQvAQu-W4YOUFH1gnr-ZS3rg_GNkRhI8rQ-I,16761
+brainstate/nn/_dynamics/_state_delay.py,sha256=KfsXXoFg8hv9wZgY-p5_-cTGbY88rviJmCfnGiDgboo,16905
 brainstate/nn/_dynamics/_synouts.py,sha256=9TGAc-nVa50th7KKn4oKLbro-4W4rwxYvp-eu7ksAIE,4491
 brainstate/nn/_dynamics/_synouts_test.py,sha256=V_jDswRN4VvEXD-2yJO3VA1TALgX0HK6oPBQiUntOWc,2266
 brainstate/nn/_elementwise/__init__.py,sha256=PK8oq1K_EG2941AiUyLxCWoRdWvMO3yt8ZJbw3Lkhu8,935
@@ -137,8 +137,8 @@ brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7
 brainstate/util/_struct.py,sha256=0exv0oOiSt1hmx20Y4J2-pCGtCTx13WcAlEYSBkyung,17640
 brainstate/util/_tracers.py,sha256=0r5T4nhxMzI79NtqroqitsdMT4YfpgV5RdYJLS5uJ0w,2285
 brainstate/util/_visualization.py,sha256=n4ZVz10z7VBqA0cKO6vyHwEMprWJgPeEqtITzDMai2Y,1519
-brainstate-0.1.0.post20241220.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20241220.dist-info/METADATA,sha256=lfkUbD1vYx4bikkUolrL_pCp4zryfvGExXyWP_QH5tM,3533
-brainstate-0.1.0.post20241220.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20241220.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20241220.dist-info/RECORD,,
+brainstate-0.1.0.post20241221.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20241221.dist-info/METADATA,sha256=amK3qX-uzU9N38oHF7v7CZ-xyyviTzvHivKvBU0Lths,3533
+brainstate-0.1.0.post20241221.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20241221.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20241221.dist-info/RECORD,,