brainstate 0.1.2__py2.py3-none-any.whl → 0.1.4__py2.py3-none-any.whl

{brainstate-0.1.2.dist-info → brainstate-0.1.4.dist-info}/RECORD

@@ -1,15 +1,15 @@
-brainstate/__init__.py,sha256=50znZ0wsJBdgX9Wvlv4xvsuNGp4XdgR60LTFrplmBz0,1496
-brainstate/_compatible_import.py,sha256=96CR40jbcrvaBXO0Hrlm7Cpde-4nRj-6z-P753bitFc,5001
+brainstate/__init__.py,sha256=A_OD4cJiVu3xpthNGJh6fhRjCKsI7_Mxsow3Al4m2-w,1496
+brainstate/_compatible_import.py,sha256=LUSZlA0APWozxM8Kf9pZrM2YbwY7X3jzVVHZInaBL7Y,4630
 brainstate/_state.py,sha256=o5Kk4bGwVz6Dfj9dlmZqdh6zUXcz6Tvc6WOjH9ajlIU,60716
 brainstate/_state_test.py,sha256=b6uvZdVRyC4n6-fYzmHNry1b-gJ6zE_kRSxGinqiHaw,1638
 brainstate/_utils.py,sha256=j-b239RHfC5BnvhGbSExQpdY21LrMEyWMSHBdNGThOI,1657
 brainstate/environ.py,sha256=VgtG0S_aR1g_1gplRWg_v2ZrcS-F6LZk35BPCBgsIIA,17660
 brainstate/environ_test.py,sha256=QD6sPCKNtqemVCGwkdImjMazatrvvLr6YeAVcfUnVVY,2045
-brainstate/mixin.py,sha256=1PPifp91eKEeMJO8cLjZPmXehPBE5tIoRhMVpvysquo,11272
+brainstate/mixin.py,sha256=LOS8WPfB6OBevdZ9wh_U6qzuw4mrCNAMcJy0yIitIBg,11522
 brainstate/mixin_test.py,sha256=Oq_0fwC9vpXDN4t4dTBhWzLdFDNlcYsrcip14F1yECI,3079
-brainstate/surrogate.py,sha256=Iu1-WWJKH5CIjXIHKwEF8ufxVYpFZCltU_6dEH9DYqk,53470
+brainstate/surrogate.py,sha256=yPc-uhJl1r14leoKAAgTWlqBZoeH-uYUAeizuQLkJec,53455
 brainstate/transform.py,sha256=OvshYpPnp3YXPG6riy15Ve7jcX8t0aaATl1xZnsFeic,858
-brainstate/typing.py,sha256=XEOzcvMos8jcVeOf2rZb20M4A89xUEUpEIHu9qit2L8,10485
+brainstate/typing.py,sha256=cIlcgaWbvsFyEb5GgwZD7NHJh9qAp8l-N-rYN5W9FsU,10519
 brainstate/augment/__init__.py,sha256=Q9-JIwQ1FNn8VLS1MA9MrSylbvUjWSw98whrI3NIuKo,1229
 brainstate/augment/_autograd.py,sha256=zMMLZOidQq2p96wzEOgR-MynV2JH1l1AnoJ28eVwD-s,30100
 brainstate/augment/_autograd_test.py,sha256=UhNd41luca_Kj9a8byL3Dq64Ta55WU-je-LCiK_z0Vc,45060
@@ -25,14 +25,14 @@ brainstate/compile/_conditions.py,sha256=fFSJWRFZxriAZFASj2lvlXpxbH323XcfsoB9ni2
 brainstate/compile/_conditions_test.py,sha256=bh8_D49Fl-s_xFYkO4joZJPgFGhw_vd6Wo8QB1HDKwE,8640
 brainstate/compile/_error_if.py,sha256=l8flahWEui9-9IavtZQjyutIeR8DU0jieahRvAXUxuw,2651
 brainstate/compile/_error_if_test.py,sha256=PjTIQOflegqnqJb9kO5YqwjW1OfrBmAtWlY_xn8rtDY,1897
-brainstate/compile/_jit.py,sha256=TN6Fk7qTHpZ2rzfBf4uyL9oFoU_IEpaDQhn1R_d_bd4,13858
+brainstate/compile/_jit.py,sha256=rpM0P2t7pKyE2ak-J-TIHL4LH9uu-IsuXnklWoE1NBA,14496
 brainstate/compile/_jit_test.py,sha256=zD7kck9SQJGmUDolh9P4luKwQ21fBGje1Z4STTEXIuA,4135
 brainstate/compile/_loop_collect_return.py,sha256=qOYGoD2TMZd2Px6y241GWNwSjOLyaKlvzOwJWIPAbxg,23341
 brainstate/compile/_loop_collect_return_test.py,sha256=pEJdcOthEM17q5kXYhgR6JfzzqibijD_O1VPXVe_Ml4,1808
 brainstate/compile/_loop_no_collection.py,sha256=tPkSxee41VexWEALpN1uuT78BDrX3uT1FHv8WZtov4c,7549
 brainstate/compile/_loop_no_collection_test.py,sha256=ivavF59xep_g9bV1SSdXd5E1f6nhc7EUBfcbgHpxbfg,1419
-brainstate/compile/_make_jaxpr.py,sha256=nW5VV9fcqD6bo0GBWxGT2005ANox1ym7ChhGV_q5Aig,33463
-brainstate/compile/_make_jaxpr_test.py,sha256=wvZHl7ecQrmD6C_Nxy8z64l6r3Tj1kQEM54UNIU1Er4,4416
+brainstate/compile/_make_jaxpr.py,sha256=fhaazC2Lj8-Gk3VpWLSytzVbP9dLkuTAn_1Y0SBtC9s,35864
+brainstate/compile/_make_jaxpr_test.py,sha256=xPusEJikMQRfoOmUFrugGcE5UpRm0giHoL_NPomN5rI,4791
 brainstate/compile/_progress_bar.py,sha256=iTp-CkcK5aMbiA0ZdlSfEXXrc4_anifYhyDzl19dksI,7488
 brainstate/compile/_unvmap.py,sha256=9S42MeTmFJa8nfBI_AjEfrAdUsDmM3KFat59O8zXIEw,4120
 brainstate/compile/_util.py,sha256=QD7lvS4Zb3P2HPNIm6mG8Rl_Lk2tQj0znOCPcH95XnI,6275
@@ -43,9 +43,9 @@ brainstate/functional/_normalization.py,sha256=_F5u9iJm8KJfxlnmV_LwJXtrAZOJj357a
 brainstate/functional/_others.py,sha256=iAqvLCmXQz6C80EjE0-5bwXCBKHoKqsggneiOAXNmQo,1699
 brainstate/functional/_spikes.py,sha256=YJ4pa5LlywximSE6ZAAs9g2k-8cO95mwxdD29HhT-oM,7051
 brainstate/graph/__init__.py,sha256=noo4TjBg6iEhjjwk0sAGUhR7Ge-z8Vnc2rLYUvnqttw,1295
-brainstate/graph/_graph_node.py,sha256=zxtfzv6RXs9ylTerW5scUZ4mHic66df5WTHCrUzHLDA,6836
+brainstate/graph/_graph_node.py,sha256=kl0CDKIcbXd9YzrZIKDhDeT8bboaoR5_OPwtL90eUEo,6835
 brainstate/graph/_graph_node_test.py,sha256=g6nrCCI-awfxhZLnjmYhfQvynumWJjz7fPxgIF-VSr4,2832
-brainstate/graph/_graph_operation.py,sha256=gnLZB5YLgsnwVAZugv_5sryzU9_QYyz2r0-D5y1ZG7Q,63405
+brainstate/graph/_graph_operation.py,sha256=h2rdydWr9AHBMVlyYpBN9GdbVkPIgyjsctZXa8EmxQI,63401
 brainstate/graph/_graph_operation_test.py,sha256=nxMwEoaRZpTkIeR3h_Aduuys0Bm1DRlYLd_CXkErTSg,19715
 brainstate/init/__init__.py,sha256=R1dHgub47o-WJM9QkFLc7x_Q7GsyaKKDtrRHTFPpC5g,1097
 brainstate/init/_base.py,sha256=A5VIJV_xPyCD4Lo4RXe2WPlWvADAy_UPhexuwzp-6EI,1633
@@ -54,53 +54,51 @@ brainstate/init/_random_inits.py,sha256=GYovk2h2Lc1jY5ByW4R8OoZ71jbaE3BYGTfDw4-a
 brainstate/init/_random_inits_test.py,sha256=c1pUdWXQxL-seEvjO_8fBkaDL6T63b6ALNfQrNuuBbg,5287
 brainstate/init/_regular_inits.py,sha256=tcZDTXR8jiRqERuJleg_9JRGGzJj-V3LhTMcr_gt_k0,3149
 brainstate/init/_regular_inits_test.py,sha256=bvv0AOfLOEP0BIQIBLztKw3EPyEp7n2fHW8PUrmWuHk,1820
-brainstate/nn/__init__.py,sha256=dUEH_S5jRTq0nY1xNkpN_4MQrmHXHGSmGp3xgGenqNc,1924
+brainstate/nn/__init__.py,sha256=nDIBDlzHaBy-vk4cb7FiPO6zoNe2tI5qvHIK8O7yrOU,3721
 brainstate/nn/_collective_ops.py,sha256=rWcjqaP0rW6sXdhE0fjtDi4twmq5VvhNEg8AoiQ-tDU,17390
 brainstate/nn/_collective_ops_test.py,sha256=bwq0DApcsk0_2xpxMl0_e2cGKT63g5rSngpigCm07ps,1409
 brainstate/nn/_common.py,sha256=Pt8P1LgE0qW3QnfX8CQQGH3yZ78RL7NkIqqYEw_z8xs,7096
+brainstate/nn/_conv.py,sha256=Zk-yj34n6CkjntcM9xpMGLTxKNfWdIWsTsoGbtdL0yU,18448
+brainstate/nn/_conv_test.py,sha256=2lcUTG7twkyhuyKwuBux-NgU8NU_W4Cp1-G8EyDJ_uk,8862
+brainstate/nn/_delay.py,sha256=l36FBgNhfL64tM3VGOsJNTtKr44HjxxtBWMFFCm3Pks,17361
+brainstate/nn/_dropout.py,sha256=Dq3hQrOBT6gODlDbcoag6zLGXTU_p_2MhnmVsV57Hds,17783
+brainstate/nn/_dropout_test.py,sha256=9i2ZW5by0S9zfVHmZKLN0WpI1bXH6h9-QLXamG6lVXA,4432
+brainstate/nn/_dynamics.py,sha256=Rx4Hwl87H94DgMM644hDi3nN9D0wvXkpWrog_Ek1sr4,47734
+brainstate/nn/_dynamics_test.py,sha256=w7AV57LdhbBNYprdFpKq8MFSCbXKVkGgp_NbL3ANX3I,2769
+brainstate/nn/_elementwise.py,sha256=4czeJWGQopV49iZo8DuN_WzAbXoMC1gtqaGjlON6e7c,33291
+brainstate/nn/_elementwise_test.py,sha256=_dd9eX2ZJ7p24ahuoapCaRTZ0g1boufXMyqHFx1d4WY,5688
+brainstate/nn/_embedding.py,sha256=SaAJbgXmuJ8XlCOX9ob4yvmgh9Fk627wMguRzJMJ1H8,2138
 brainstate/nn/_exp_euler.py,sha256=ndDB43PM4jsZKu_zdLTZ2-ojnuNrg55LZap23oBTtdA,3493
 brainstate/nn/_exp_euler_test.py,sha256=XD--qMbGHrHa3WtcPMmJKk59giDcEhSqZuBOmTNYUr8,1227
-brainstate/nn/_module.py,sha256=ksbmO_-bZzaP8XTRAcN2oi2ZMdSuSv-XVeWc2VJncVc,12808
+brainstate/nn/_fixedprob.py,sha256=KGXohiU0wZnFIQDuwiRUTFsbsr8R0p8zgi5UZDuv1Bk,10004
+brainstate/nn/_fixedprob_test.py,sha256=qbRBh-MpMtEOsg492gFu2w9-FOP9z_bXapm-Q0gLLYM,3929
+brainstate/nn/_inputs.py,sha256=wPOfPE4IesNoDmxZJxqR0siBlJioEX-_1IZ2cltAIpM,20605
+brainstate/nn/_linear.py,sha256=5WuhcqU-uBUC91vnwezQYMHPKmlZPDgIJ5UpffxoX1I,14472
+brainstate/nn/_linear_mv.py,sha256=6hDXx4yPqRSa7uIsW9f9eJuy23dcXN9Mp2_lSvw8BDA,2635
+brainstate/nn/_linear_mv_test.py,sha256=ZCM1Zy6mImQfCfdZOGnTwkiLLPXK5yalv1Ts9sWZuPA,3864
+brainstate/nn/_linear_test.py,sha256=eIS-VCR3QmXB_byO1Uexg65Pv48CBRUA_Je-UGrFVTY,2925
+brainstate/nn/_ltp.py,sha256=_najNUyfaFYcOlUTm7ThJopInbos3kwJyrm-QUfI-hc,861
+brainstate/nn/_module.py,sha256=tleUiEd3ytEx4fZLu9NButELVlbKZNIQbXX78bdK8QU,12770
 brainstate/nn/_module_test.py,sha256=_RKHY8LHg6nbFHkvPAsEPVpp9rW6XfrcF8YD68u4o4w,8956
+brainstate/nn/_neuron.py,sha256=2walTScvL034LS53pArDASXz6z26SSPbmCvchWWjkUU,27441
+brainstate/nn/_neuron_test.py,sha256=QF8pixUqA5Oj7MrNi2NR8VAnfGpAvNpwV2mBc3e_pTY,6393
+brainstate/nn/_normalizations.py,sha256=4B1YVI6AZyJLHseJK9eiSeMOIq--0ZuDG58iC9mrekE,37387
+brainstate/nn/_normalizations_test.py,sha256=JTqpH265GT0eKw3hXmT0qC8ZM5NkUe98of0uYS7_5Us,2456
+brainstate/nn/_poolings.py,sha256=h5aIj_K3X2HSMjrOtkgwiCpvLxfXXEUvFFa57Cxp5bI,46969
+brainstate/nn/_poolings_test.py,sha256=qje9PVWvPGiYOv6UlTEWfpqqjpv4Xop5rO0ATcgcF0w,7497
+brainstate/nn/_projection.py,sha256=_TW-YmfPqUDkekMMK7tXrdLbjT4Ou2sUAecTNNn7FVM,17498
+brainstate/nn/_rate_rnns.py,sha256=OQHMyq9kKd2q44dOHLuTqHKTGuraPMqXVutFP_LYIyU,20676
+brainstate/nn/_rate_rnns_test.py,sha256=__hhx7e6LX_1mDLLQyIi4TNCaFAWnOVSTIgwHNjzf2g,2444
+brainstate/nn/_readout.py,sha256=OJjSba5Wr7dtUXqYhAv1D7BUGOI-lAmg6urxPBrZe3c,7116
+brainstate/nn/_readout_test.py,sha256=L2T0-SkiACxkY_I5Pbnbmy0Zw3tbpV3l5xVzAw42f2g,2136
+brainstate/nn/_stp.py,sha256=-ahDEqSp8bQsU_nUK4jks8fjMYKgIbO0v7zpyGVuXtA,8645
+brainstate/nn/_synapse.py,sha256=liKhmHZFQpEdu1R76AljeJMAKz5HK8Yq95ENa6FLVAc,20145
+brainstate/nn/_synapse_test.py,sha256=xmCWFxZUIM2YtmW5otKnADGCCK__4JpXmSYcZ3wzlQM,4994
+brainstate/nn/_synaptic_projection.py,sha256=UFgzsMB1VZ9ieumwmaTIC7irLrZ4pmwfiuOYomkwXG4,17917
+brainstate/nn/_synouts.py,sha256=jWQP1-qXFpdYgyUSJNFD7_bk4_-67ok36br-OzbcSXY,4524
+brainstate/nn/_synouts_test.py,sha256=sfjotlS--4hT22Vb5RfmpfmXABa8z-VdxZYUSocMmlU,2244
 brainstate/nn/_utils.py,sha256=epfELIy1COgdS9z5be-fmbFhagNugcIHpw4ww-HlkSY,3123
 brainstate/nn/metrics.py,sha256=TXCB_yGQzamklJCI5FGOAZ5dihtY-gjYlSi2SCLC3LA,14700
-brainstate/nn/_dyn_impl/__init__.py,sha256=J1hMqNey1921tZppCjuHApB6oc-5aGhWXyOXh96CeWM,1322
-brainstate/nn/_dyn_impl/_dynamics_neuron.py,sha256=6xsGF2Oy_oD6bcTdfwqnHBtfdAJREIl5qvNUZKZRlnQ,27482
-brainstate/nn/_dyn_impl/_dynamics_neuron_test.py,sha256=QF8pixUqA5Oj7MrNi2NR8VAnfGpAvNpwV2mBc3e_pTY,6393
-brainstate/nn/_dyn_impl/_dynamics_synapse.py,sha256=b2QMEzasbXTp7k1XWNg7nRt4-WddCym0CSEErjPRIWA,27449
-brainstate/nn/_dyn_impl/_dynamics_synapse_test.py,sha256=xmCWFxZUIM2YtmW5otKnADGCCK__4JpXmSYcZ3wzlQM,4994
-brainstate/nn/_dyn_impl/_inputs.py,sha256=Fe3tgSwJMrq2HfZJq9CTgNy24Ji6tnLCqB4hkhaY_K0,20675
-brainstate/nn/_dyn_impl/_rate_rnns.py,sha256=SXlrHZhih4LD99Jg8OwIWur5UQbd6E_zotkHvkFFstc,20715
-brainstate/nn/_dyn_impl/_rate_rnns_test.py,sha256=__hhx7e6LX_1mDLLQyIi4TNCaFAWnOVSTIgwHNjzf2g,2444
-brainstate/nn/_dyn_impl/_readout.py,sha256=UzaJ6WIMpNjqHolhjHTpyCZ5E2nSDucvuV5PBW7rOPs,7151
-brainstate/nn/_dyn_impl/_readout_test.py,sha256=L2T0-SkiACxkY_I5Pbnbmy0Zw3tbpV3l5xVzAw42f2g,2136
-brainstate/nn/_dynamics/__init__.py,sha256=j1HSWu01wf5-KjSaNhBC9utVGDALOhUsFPrLPcPPDsM,1208
-brainstate/nn/_dynamics/_dynamics_base.py,sha256=stBAkb8f973vJe66p0tve-0AO_j008BCpAA96FIPnTI,43363
-brainstate/nn/_dynamics/_dynamics_base_test.py,sha256=w7AV57LdhbBNYprdFpKq8MFSCbXKVkGgp_NbL3ANX3I,2769
-brainstate/nn/_dynamics/_projection_base.py,sha256=O86VfZkd8Qd5_8CDK-4eCGS6Jkp1CKyW6J9__bQnmo0,12893
-brainstate/nn/_dynamics/_state_delay.py,sha256=kSgel1OSFB92WcfGTsIjgb-KW389ydXTcLbTUjylU-A,16867
-brainstate/nn/_dynamics/_synouts.py,sha256=Py9DM3o-JmkoqYyQM5vl89-SBqTBRdk2CUYgYHYu3No,4453
-brainstate/nn/_dynamics/_synouts_test.py,sha256=sfjotlS--4hT22Vb5RfmpfmXABa8z-VdxZYUSocMmlU,2244
-brainstate/nn/_elementwise/__init__.py,sha256=PK8oq1K_EG2941AiUyLxCWoRdWvMO3yt8ZJbw3Lkhu8,935
-brainstate/nn/_elementwise/_dropout.py,sha256=ymr87YYA5x5C_c9HzNLGPJK_ebJ22Y-KsK-gZJY_2MU,17733
-brainstate/nn/_elementwise/_dropout_test.py,sha256=9i2ZW5by0S9zfVHmZKLN0WpI1bXH6h9-QLXamG6lVXA,4432
-brainstate/nn/_elementwise/_elementwise.py,sha256=yO6BukJFLfnICPELEFp2kfutFM8Jeg7TaERXFp-aNrk,33304
-brainstate/nn/_elementwise/_elementwise_test.py,sha256=_dd9eX2ZJ7p24ahuoapCaRTZ0g1boufXMyqHFx1d4WY,5688
-brainstate/nn/_event/__init__.py,sha256=OLteeA3d1Kq8yi3XETA-rFYtYpojZRpY1J_h5MpH3rY,919
-brainstate/nn/_event/_fixedprob_mv.py,sha256=MpVCiYGiDC8GAuJWtYXKrpzP9yDtnYHolHlwyzoQ1r4,6619
-brainstate/nn/_event/_fixedprob_mv_test.py,sha256=qbRBh-MpMtEOsg492gFu2w9-FOP9z_bXapm-Q0gLLYM,3929
-brainstate/nn/_event/_linear_mv.py,sha256=KslIdse7KZ7FWm5qcOMC-COM__1rB-sUJmXvlXY11y8,2683
-brainstate/nn/_event/_linear_mv_test.py,sha256=ZCM1Zy6mImQfCfdZOGnTwkiLLPXK5yalv1Ts9sWZuPA,3864
-brainstate/nn/_interaction/__init__.py,sha256=TTY_SeNrdx4VnUSw6vdyl02OHdS9Qs15cWBp6kjsyNQ,1289
-brainstate/nn/_interaction/_conv.py,sha256=QMZR4HRAYe8Nzr3qExp0PWr6wuUp2rh5YxmPq1r1vcY,18461
-brainstate/nn/_interaction/_conv_test.py,sha256=2lcUTG7twkyhuyKwuBux-NgU8NU_W4Cp1-G8EyDJ_uk,8862
-brainstate/nn/_interaction/_embedding.py,sha256=8HC-J3RhvGGrvqA0-BBaFDk6tmlunPP6bNohjRSQzYo,2151
-brainstate/nn/_interaction/_linear.py,sha256=jk765kASgwVG1r-xjoM49Thfa1Cu32mscy7ci79bjDI,14729
-brainstate/nn/_interaction/_linear_test.py,sha256=eIS-VCR3QmXB_byO1Uexg65Pv48CBRUA_Je-UGrFVTY,2925
-brainstate/nn/_interaction/_normalizations.py,sha256=4yoU04wK21_zk6Egpwb9YN_pSWUH5Qa6hJ1bLxX7wQw,37400
-brainstate/nn/_interaction/_normalizations_test.py,sha256=JTqpH265GT0eKw3hXmT0qC8ZM5NkUe98of0uYS7_5Us,2456
-brainstate/nn/_interaction/_poolings.py,sha256=2hJDACuHv3CkO-nyyAqVZp7Sa3UFa_ROFGiYmhXWqKk,46982
-brainstate/nn/_interaction/_poolings_test.py,sha256=7-jEs_UcMNWm8q6NbrRRX3bTn14jBjGiAX0KBaROng0,7510
 brainstate/optim/__init__.py,sha256=7Ao0LCtDNAoxSRSXiLLKnd1_4mR2GSExizpN38il-Fo,1195
 brainstate/optim/_base.py,sha256=P37k8w46iQZZZnFLa5OF83Sb8DLQfYMez9ZRObBqVsE,1835
 brainstate/optim/_lr_scheduler.py,sha256=0t7kl35MLmc0UcSu587nI28cI3SohYhy_19Al-5dNTM,15289
@@ -115,19 +113,19 @@ brainstate/random/_rand_seed.py,sha256=1ZdfFZWyOhpd72EDdEDmpkp3yoLVwdv-sGI9BwiZf
 brainstate/random/_rand_seed_test.py,sha256=waXXfch57X1XE1zDnCRokT6ziZOK0g-lYE80o6epDYM,1536
 brainstate/random/_rand_state.py,sha256=dTO7wqmuTYRdPy7ItsrK-7aNt5QQXOsZ4XwiH7mzmy8,55170
 brainstate/random/_random_for_unit.py,sha256=kGp4EUX19MXJ9Govoivbg8N0bddqOldKEI2h_TbdONY,2057
-brainstate/util/__init__.py,sha256=zwdw0253tY2kgWHHPPwaU6JHvjw4rioUnRmSAWF4Po0,1569
-brainstate/util/_caller.py,sha256=RBRwu1HT_Lww7EE3WTjzrwwKqVWg7dn_JXKmO9ojM38,2752
-brainstate/util/_error.py,sha256=GBLaU_glQlgdrjBYIMiO3Dy5N_ST6-QNsgAOa9kjUjw,1757
-brainstate/util/_others.py,sha256=vGo4uW0vd7wL5DYtxKWqXhoNmI-O_p26HsucoMnhdVA,16407
-brainstate/util/_pretty_pytree.py,sha256=PHMiES2vN_bke_GeHaSYHQZUtXtR-HkcMch8HQJ8Cw8,33536
+brainstate/util/__init__.py,sha256=6efwr63osmqviNU_6_Nufag19PwxRDFvDAZrq6sH5yo,1555
 brainstate/util/_pretty_pytree_test.py,sha256=Dn0TdjX6wLBXaTD4jfYTu6cKfFHwKSxi4_3bX7kB_IA,5621
-brainstate/util/_pretty_repr.py,sha256=7Xp7IFNUeP7cGlpvwwJyBslbQVnXEqC1I6neV1Jx1S8,10527
-brainstate/util/_pretty_table.py,sha256=uJVaamFGQ4nKP8TkEGPWXHpzjMecDo2q1Ah6XtRjdPY,108117
-brainstate/util/_scaling.py,sha256=U6DM-afPrLejiGqo1Nla7z4YbTBVicctsBEweurr_mk,7524
-brainstate/util/_struct.py,sha256=7HbbQNrZ3zxYw93MU1bUZ9ZPBKftYOVKuXEochLSErw,17479
+brainstate/util/caller.py,sha256=RBRwu1HT_Lww7EE3WTjzrwwKqVWg7dn_JXKmO9ojM38,2752
+brainstate/util/error.py,sha256=GBLaU_glQlgdrjBYIMiO3Dy5N_ST6-QNsgAOa9kjUjw,1757
 brainstate/util/filter.py,sha256=skuq51y0P5-N611fQ-zB2QyaJIFqN_WbROzZHv0RgOc,14086
-brainstate-0.1.2.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.2.dist-info/METADATA,sha256=Mv44muheQeXJ67RF5GQrkSxVWfj4VAZnkse6_7z5OTY,4135
-brainstate-0.1.2.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.2.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.2.dist-info/RECORD,,
+brainstate/util/others.py,sha256=vGo4uW0vd7wL5DYtxKWqXhoNmI-O_p26HsucoMnhdVA,16407
+brainstate/util/pretty_pytree.py,sha256=HmXZpKGps1Rs0DEN7r1eYN6xMM6FgDpCC8SdEHjOsvo,33534
+brainstate/util/pretty_repr.py,sha256=7Xp7IFNUeP7cGlpvwwJyBslbQVnXEqC1I6neV1Jx1S8,10527
+brainstate/util/pretty_table.py,sha256=uJVaamFGQ4nKP8TkEGPWXHpzjMecDo2q1Ah6XtRjdPY,108117
+brainstate/util/scaling.py,sha256=U6DM-afPrLejiGqo1Nla7z4YbTBVicctsBEweurr_mk,7524
+brainstate/util/struct.py,sha256=7HbbQNrZ3zxYw93MU1bUZ9ZPBKftYOVKuXEochLSErw,17479
+brainstate-0.1.4.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.4.dist-info/METADATA,sha256=EWarKpuYFIN0ITFn8lwsUy0h9-KjrVPDwcvjQMHFRto,4135
+brainstate-0.1.4.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.4.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.4.dist-info/RECORD,,

brainstate/nn/_dyn_impl/__init__.py

@@ -1,42 +0,0 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-
-from ._dynamics_neuron import *
-from ._dynamics_neuron import __all__ as dyn_neuron_all
-from ._dynamics_synapse import *
-from ._dynamics_synapse import __all__ as dyn_synapse_all
-from ._inputs import *
-from ._inputs import __all__ as inputs_all
-from ._rate_rnns import *
-from ._rate_rnns import __all__ as rate_rnns
-from ._readout import *
-from ._readout import __all__ as readout_all
-
-__all__ = (
-    dyn_neuron_all
-    + dyn_synapse_all
-    + inputs_all
-    + rate_rnns
-    + readout_all
-)
-
-del (
-    dyn_neuron_all,
-    dyn_synapse_all,
-    inputs_all,
-    readout_all,
-    rate_rnns,
-)

brainstate/nn/_dynamics/__init__.py

@@ -1,37 +0,0 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-from ._dynamics_base import *
-from ._dynamics_base import __all__ as dyn_all
-from ._projection_base import *
-from ._projection_base import __all__ as projection_all
-from ._state_delay import *
-from ._state_delay import __all__ as state_delay_all
-from ._synouts import *
-from ._synouts import __all__ as synouts_all
-
-__all__ = (
-    dyn_all
-    + projection_all
-    + state_delay_all
-    + synouts_all
-)
-
-del (
-    dyn_all,
-    projection_all,
-    state_delay_all,
-    synouts_all
-)

brainstate/nn/_dynamics/_projection_base.py

@@ -1,362 +0,0 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-from typing import Union, Callable, Optional
-
-from brainstate._state import State
-from brainstate.mixin import AlignPost, ParamDescriber, BindCondData, JointTypes
-from brainstate.nn._collective_ops import call_order
-from brainstate.nn._module import Module
-from brainstate.util._others import get_unique_name
-from ._dynamics_base import Dynamics, maybe_init_prefetch, Prefetch, PrefetchDelayAt
-from ._synouts import SynOut
-
-__all__ = [
-    'AlignPostProj',
-    'DeltaProj',
-    'CurrentProj',
-]
-
-
-class Interaction(Module):
-    __module__ = 'brainstate.nn'
-
-
-def _check_modules(*modules):
-    # checking modules
-    for module in modules:
-        if not callable(module) and not isinstance(module, State):
-            raise TypeError(
-                f'The module should be a callable function or a brainstate.State, but got {module}.'
-            )
-    return tuple(modules)
-
-
-def call_module(module, *args, **kwargs):
-    if callable(module):
-        return module(*args, **kwargs)
-    elif isinstance(module, State):
-        return module.value
-    else:
-        raise TypeError(
-            f'The module should be a callable function or a brainstate.State, but got {module}.'
-        )
-
-
-def is_instance(x, cls) -> bool:
-    return isinstance(x, cls)
-
-
-def get_post_repr(label, syn, out):
-    if label is None:
-        return f'{syn.identifier} // {out.identifier}'
-    else:
-        return f'{label}{syn.identifier} // {out.identifier}'
-
-
-def align_post_add_bef_update(
-    syn_desc: ParamDescriber[AlignPost],
-    out_desc: ParamDescriber[BindCondData],
-    post: Dynamics,
-    proj_name: str,
-    label: str,
-):
-    # synapse and output initialization
-    _post_repr = get_post_repr(label, syn_desc, out_desc)
-    if not post._has_before_update(_post_repr):
-        syn_cls = syn_desc()
-        out_cls = out_desc()
-
-        # synapse and output initialization
-        post.add_current_input(proj_name, out_cls, label=label)
-        post._add_before_update(_post_repr, _AlignPost(syn_cls, out_cls))
-    syn = post._get_before_update(_post_repr).syn
-    out = post._get_before_update(_post_repr).out
-    return syn, out
-
-
-class _AlignPost(Module):
-    def __init__(
-        self,
-        syn: Dynamics,
-        out: BindCondData
-    ):
-        super().__init__()
-        self.syn = syn
-        self.out = out
-
-    def update(self, *args, **kwargs):
-        self.out.bind_cond(self.syn(*args, **kwargs))
-
-
-class AlignPostProj(Interaction):
-    """
-    Align-post projection of the neural network.
-
-
-    Examples
-    --------
-
-    Here is an example of using the `AlignPostProj` to create a synaptic projection.
-    Note that this projection needs the manual input of pre-synaptic spikes.
-
-    >>> import brainstate
-    >>> import brainevent
-    >>> import brainunit as u
-    >>> n_exc = 3200
-    >>> n_inh = 800
-    >>> num = n_exc + n_inh
-    >>> pop = brainstate.nn.LIFRef(
-    ...        num,
-    ...        V_rest=-49. * u.mV, V_th=-50. * u.mV, V_reset=-60. * u.mV,
-    ...        tau=20. * u.ms, tau_ref=5. * u.ms,
-    ...        V_initializer=brainstate.init.Normal(-55., 2., unit=u.mV)
-    ... )
-    >>> pop.reset_state()
-    >>> E = brainstate.nn.AlignPostProj(
-    ...        comm=brainevent.nn.FixedProb(n_exc, num, prob=80 / num, weight=1.62 * u.mS),
-    ...        syn=brainstate.nn.Expon.desc(num, tau=5. * u.ms),
-    ...        out=brainstate.nn.CUBA.desc(scale=u.volt),
-    ...        post=pop
-    ... )
-    >>> exe_current = E(pop.spike.value)
-
-
-
-    """
-    __module__ = 'brainstate.nn'
-
-    def __init__(
-        self,
-        *modules,
-        comm: Callable,
-        syn: Union[ParamDescriber[AlignPost], AlignPost],
-        out: Union[ParamDescriber[SynOut], SynOut],
-        post: Dynamics,
-        label: Optional[str] = None,
-    ):
-        super().__init__(name=get_unique_name(self.__class__.__name__))
-
-        # checking modules
-        self.modules = _check_modules(*modules)
-
-        # checking communication model
-        if not callable(comm):
-            raise TypeError(
-                f'The communication should be an instance of callable function, but got {comm}.'
-            )
-
-        # checking synapse and output models
-        if is_instance(syn, ParamDescriber[AlignPost]):
-            if not is_instance(out, ParamDescriber[SynOut]):
-                if is_instance(out, ParamDescriber):
-                    raise TypeError(
-                        f'The output should be an instance of describer {ParamDescriber[SynOut]} when '
-                        f'the synapse is an instance of {AlignPost}, but got {out}.'
-                    )
-                raise TypeError(
-                    f'The output should be an instance of describer {ParamDescriber[SynOut]} when '
-                    f'the synapse is a describer, but we got {out}.'
-                )
-            merging = True
-        else:
-            if is_instance(syn, ParamDescriber):
-                raise TypeError(
-                    f'The synapse should be an instance of describer {ParamDescriber[AlignPost]}, but got {syn}.'
-                )
-            if not is_instance(out, SynOut):
-                raise TypeError(
-                    f'The output should be an instance of {SynOut} when the synapse is '
-                    f'not a describer, but we got {out}.'
-                )
-            merging = False
-        self.merging = merging
-
-        # checking post model
-        if not is_instance(post, Dynamics):
-            raise TypeError(
-                f'The post should be an instance of {Dynamics}, but got {post}.'
-            )
-
-        if merging:
-            # synapse and output initialization
-            syn, out = align_post_add_bef_update(syn_desc=syn,
-                                                 out_desc=out,
-                                                 post=post,
-                                                 proj_name=self.name,
-                                                 label=label)
-        else:
-            post.add_current_input(self.name, out)
-
-        # references
-        self.comm = comm
-        self.syn: JointTypes[Dynamics, AlignPost] = syn
-        self.out: BindCondData = out
-        self.post: Dynamics = post
-
-    @call_order(2)
-    def init_state(self, *args, **kwargs):
-        for module in self.modules:
-            maybe_init_prefetch(module, *args, **kwargs)
-
-    def update(self, *args):
-        # call all modules
-        for module in self.modules:
-            x = call_module(module, *args)
-            args = (x,)
-        # communication module
-        x = self.comm(*args)
-        # add synapse input
-        self.syn.add_delta_input(self.name, x)
-        if not self.merging:
-            # synapse and output interaction
-            conductance = self.syn()
-            self.out.bind_cond(conductance)
-
-
-class DeltaProj(Interaction):
-    r"""Full-chain of the synaptic projection for the Delta synapse model.
-
-    The synaptic projection requires the input is the spiking data, otherwise
-    the synapse is not the Delta synapse model.
-
-    The ``full-chain`` means that the model needs to provide all information needed for a projection,
-    including ``pre`` -> ``delay`` -> ``comm`` -> ``post``.
-
-    **Model Descriptions**
-
-    .. math::
-
-        I_{syn} (t) = \sum_{j\in C} g_{\mathrm{max}} * \delta(t-t_j-D)
-
-    where :math:`g_{\mathrm{max}}` denotes the chemical synaptic strength,
-    :math:`t_j` the spiking moment of the presynaptic neuron :math:`j`,
-    :math:`C` the set of neurons connected to the post-synaptic neuron,
-    and :math:`D` the transmission delay of chemical synapses.
-    For simplicity, the rise and decay phases of post-synaptic currents are
-    omitted in this model.
-
-    #     brainstate.nn.DeltaInteraction(
-    #       LIF().prefetch('V'), bst.surrogate.ReluGrad(), comm, post
-    #     )
-
-    Args:
-      pre: The pre-synaptic neuron group.
-      delay: The synaptic delay.
-      comm: DynamicalSystem. The synaptic communication.
-      post: DynamicalSystem. The post-synaptic neuron group.
-    """
-
-    __module__ = 'brainstate.nn'
-
-    def __init__(self, *modules, comm: Callable, post: Dynamics, label=None):
-        super().__init__(name=get_unique_name(self.__class__.__name__))
-
-        self.label = label
-
-        # checking modules
-        self.modules = _check_modules(*modules)
-
-        # checking communication model
-        if not callable(comm):
-            raise TypeError(
-                f'The communication should be an instance of callable function, but got {comm}.'
-            )
-        self.comm = comm
-
-        # post model
-        if not isinstance(post, Dynamics):
-            raise TypeError(
-                f'The post should be an instance of {Dynamics}, but got {post}.'
-            )
-        self.post = post
-
-    @call_order(2)
-    def init_state(self, *args, **kwargs):
-        for module in self.modules:
-            maybe_init_prefetch(module, *args, **kwargs)
-
-    def update(self, *x):
-        for module in self.modules:
-            x = (call_module(module, *x),)
-        assert len(x) == 1, f'The output of the modules should be a single value, but got {x}.'
-        x = self.comm(x[0])
-        self.post.add_delta_input(self.name, x, label=self.label)
-
-
-class CurrentProj(Interaction):
-    """
-    Full-chain synaptic projection with the align-pre reduction and delay+synapse updating and merging.
-
-    The ``full-chain`` means that the model needs to provide all information needed for a projection,
-    including ``pre`` -> ``delay`` -> ``syn`` -> ``comm`` -> ``out`` -> ``post``.
-    Note here, compared to ``FullProjAlignPreSD``, the ``delay`` and ``syn`` are exchanged.
-
-    The ``align-pre`` means that the synaptic variables have the same dimension as the pre-synaptic neuron group.
-
-    The ``delay+synapse updating`` means that the projection first delivers the pre neuron output (usually the
-    spiking)  to the delay model, then computes the synapse states, and finally computes the synaptic current.
-
-    The ``merging`` means that the same delay model is shared by all synapses, and the synapse model with same
-    parameters (such like time constants) will also share the same synaptic variables.
-
-    Neither ``FullProjAlignPreDS`` nor ``FullProjAlignPreSD`` facilitates the event-driven computation.
-    This is because the ``comm`` is computed after the synapse state, which is a floating-point number, rather
-    than the spiking. To facilitate the event-driven computation, please use align post projections.
-
-    Args:
-      prefetch: The synaptic dynamics.
-      comm: The synaptic communication.
-      out: The synaptic output.
-      post: The post-synaptic neuron group.
-    """
-    __module__ = 'brainstate.nn'
-
-    def __init__(
-        self,
-        prefetch: Union[Prefetch, PrefetchDelayAt],
-        comm: Callable,
-        out: SynOut,
-        post: Dynamics,
-    ):
-        super().__init__(name=get_unique_name(self.__class__.__name__))
-
-        # pre-synaptic neuron group
-        if not isinstance(prefetch, (Prefetch, PrefetchDelayAt)):
-            raise TypeError(f'The pre should be a Prefetch or PrefetchDelayAt, but got {prefetch}.')
-        self.prefetch = prefetch
-
-        # check out
-        if not isinstance(out, SynOut):
-            raise TypeError(f'The out should be a SynOut, but got {out}.')
-        self.out = out
-
-        # check post
-        if not isinstance(post, Dynamics):
-            raise TypeError(f'The post should be a Dynamics, but got {post}.')
-        self.post = post
-        post.add_current_input(self.name, out)
-
-        # output initialization
-        self.comm = comm
-
-    @call_order(2)
-    def init_state(self, *args, **kwargs):
-        maybe_init_prefetch(self.prefetch, *args, **kwargs)
-
-    def update(self, *x):
-        x = self.prefetch(*x)
-        x = self.comm(x)
-        self.out.bind_cond(x)

brainstate/nn/_elementwise/__init__.py

@@ -1,22 +0,0 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-from ._dropout import *
-from ._dropout import __all__ as dropout_all
-from ._elementwise import *
-from ._elementwise import __all__ as elementwise_all
-
-__all__ = dropout_all + elementwise_all
-del dropout_all, elementwise_all