brainstate 0.1.0__py2.py3-none-any.whl → 0.1.0.post20241125__py2.py3-none-any.whl

{brainstate-0.1.0.dist-info → brainstate-0.1.0.post20241125.dist-info}/RECORD

@@ -1,5 +1,7 @@
+benchmark/COBA_2005.py,sha256=Q8PsZ0lxu14jsF3bCtlZW35iQB8S2_oFEUYQzK2hPiA,5561
+benchmark/CUBA_2005.py,sha256=_W94yOMh2ueqblU4ItEPeTLwHF0_lbEWlVNEBy0Tix0,6222
 brainstate/__init__.py,sha256=r7C3eLTg8LEusoH6PGgBFFt4ZgbketYLoLA0lQhUCsE,2098
-brainstate/_state.py,sha256=L-wDWR04ON1VAqBDv0cHdGM-DVK9lHdH1p2LCKMWoJA,28802
+brainstate/_state.py,sha256=0h3B32130Tvv41eeolYTbEGT9FZ9WUZ9yYoGyVnte_c,28808
 brainstate/_state_test.py,sha256=1boTp1w8DiCFLsPwNtlLrlIqGRpkasAmLid5bv2fgP4,2223
 brainstate/_utils.py,sha256=uJ6WWKq3yb05ZdktCQGLWOXsOJveL1H9pR7eev70Jes,1693
 brainstate/environ.py,sha256=G6r_rqfbofRbjFFalRu_DHaL7ruFTeLRXBQDXM6P-tQ,17477
@@ -10,8 +12,8 @@ brainstate/surrogate.py,sha256=YaY6RJ6kzpuPXWFjaWsxWt2MzJfdm5v_jeOR8V_jPoU,48369
 brainstate/transform.py,sha256=uryXL2vAD-oApUNvL4n6uzLPRwXXARCJ6lRwco_Gh_Y,772
 brainstate/typing.py,sha256=XaqpkfkMxBYkze5DoCQ48GSxuyO0C8rGOKeoN9qeZNo,10425
 brainstate/augment/__init__.py,sha256=BtXIBel7GbttmfBX6grxOxl0IiOJxLEa7qCGAXumamE,1286
-brainstate/augment/_autograd.py,sha256=4pbiAYCwnBuD5xpRo50HU5DvpZmnaB1Z2L-1TVY4OAE,25357
-brainstate/augment/_autograd_test.py,sha256=WIgUCENzW7pdtsR0PCHjOJnAz4eI1Ha9wlQsS6V8Yjc,41283
+brainstate/augment/_autograd.py,sha256=o9ivoEY7BmtdM1XmzdMmeRXpj6Tvn5xNB8LSGp2HKC8,25238
+brainstate/augment/_autograd_test.py,sha256=S2eEgrwTzdSi3u2nKE3u37WSThosLwx1WCP9ptJAGKo,44060
 brainstate/augment/_eval_shape.py,sha256=dGlRVHOAZ9LSRZsFi1erxgEWHrnhBO3Kq3WW11-Hvng,3819
 brainstate/augment/_eval_shape_test.py,sha256=1nnxbU7hPRbZPQWNWbQ518pw-H7FGDKKnQpZGBY9uRI,1390
 brainstate/augment/_mapping.py,sha256=cpxzVGCEYnP5jPqrowYoPXciw_-QR2F3wggrRj1OCPc,21850
@@ -35,14 +37,20 @@ brainstate/compile/_make_jaxpr_test.py,sha256=qJUtkyj50JQ6f4UJbOLhvRdkbNn3NSKibF
 brainstate/compile/_progress_bar.py,sha256=LML4DjrLSIeGYJWLjqy6BnHSz03fu1gnjf-7kljP384,3824
 brainstate/compile/_unvmap.py,sha256=ewbLLNXiI_dBsEBaVzSS0BEXNol22sd9gMzk606lSkM,4139
 brainstate/compile/_util.py,sha256=aCvkTV--g4NsqcodTdBAISt4EwgezCbKzNUV58n-Q_Y,6304
-brainstate/event/__init__.py,sha256=83a1IWZ_Oma1XhifEyt_l7jnEmRB-Q4Vd-t0PIolGkQ,1058
-brainstate/event/_csr.py,sha256=rEBPLnGCmf8kslB0UFsOnMjeg8s1oOt-a6mJrii3b_M,10467
-brainstate/event/_csr_test.py,sha256=qx3TZKnvIC5grSubDFRp7vubu2RJGeCDZFPmof-bRiA,3874
-brainstate/event/_fixed_probability.py,sha256=kZYghzEopb2lQPO3RnCxtta80hfVM-9sdSJ6gzKlcjM,9656
-brainstate/event/_fixed_probability_test.py,sha256=UFeb5ey2mLFaIPvgwrEE2HUjAIRwb3MuHsnnlfqNs9w,4119
-brainstate/event/_linear.py,sha256=kgSA9xzIMzUAl_ztKHuDG5gUZsBCr1AgvuHkAdkEyH4,7006
-brainstate/event/_linear_test.py,sha256=AKttUDxlFvdeN8j3EKSGDWaB629ANYRcDg2G_n5Tbkc,3577
+brainstate/event/__init__.py,sha256=wOBkq7kDg90M8Y9FuoXRlSEuu1ZzbIhCJ1dHeLqN6_Q,1194
+brainstate/event/_csr_benchmark.py,sha256=xrj2DSWzw0pUHAE1jRBeSRhMW7ogXvDHEdeaZGioNE4,702
+brainstate/event/_csr_mv.py,sha256=4PVSK6QuuDK6pfA5SXvUU4Cxwkd-alGJ2u7VI0irwaQ,10718
+brainstate/event/_csr_mv_benchmark.py,sha256=xrj2DSWzw0pUHAE1jRBeSRhMW7ogXvDHEdeaZGioNE4,702
+brainstate/event/_csr_mv_test.py,sha256=qx3TZKnvIC5grSubDFRp7vubu2RJGeCDZFPmof-bRiA,3874
+brainstate/event/_fixedprob_mv.py,sha256=uFJBlE3hh_QHl3TvNy061xMuLCykuAJGOP3kc-YfH2w,25187
+brainstate/event/_fixedprob_mv_benchmark.py,sha256=_F_8fH5MNMJZHeSqnq9DYMI9OgYr6JIxBKjbsgeWRv4,4720
+brainstate/event/_fixedprob_mv_test.py,sha256=jijrtJ5fnwcLmA7Tjd3vDlzwfbftmLoVTN4-MPuogVc,4201
+brainstate/event/_linear_mv.py,sha256=n1NxBUQHWPHBp_xjwpLvtKW8RMvjR5cRjdV_fh_2gGo,11788
+brainstate/event/_linear_mv_benckmark.py,sha256=hu0WqYMIa3jMoH7Fq9dgxcBjjXGFhghPx9vztyCo1KY,2411
+brainstate/event/_linear_mv_test.py,sha256=V9w41ZP2vu95CyCdCkm-j9Eftqs2kqmeBY809N1-syY,3736
 brainstate/event/_misc.py,sha256=8IpPooXjF2m0-tuo3pGHqThq2yLSNmYziy_zdurZ3NI,1040
+brainstate/event/_xla_custom_op.py,sha256=QB4jz_fUEPF-efJCVKAxwx8U79AqdcKoEg2QrGwot8I,10864
+brainstate/event/_xla_custom_op_test.py,sha256=rnkGMleXzLfJj4y5QqwfBvCCLTAHe_uabwBDniY-URM,1745
 brainstate/functional/__init__.py,sha256=j6-3Er4fgqWpvntzYCZVB3e5hoz-Z3aqvapITCuDri0,1107
 brainstate/functional/_activations.py,sha256=tsCGoV_35IyZ0mM6_4fV9v0-Vj3V9Qm55U8wSpZ_E4o,22180
 brainstate/functional/_activations_test.py,sha256=T___RlSrIfXwlkw8dg5A9EZMTZGDzv3a2evUwq_nYFg,13034
@@ -66,7 +74,7 @@ brainstate/init/_regular_inits.py,sha256=DmVMajugfyYFNUMzgFdDKMvbBu9hMWxkfDd-50u
 brainstate/init/_regular_inits_test.py,sha256=tJl4aOkclllJIfKzJTbc0cfYCw2SoBsx8_G123RnqbU,1842
 brainstate/nn/__init__.py,sha256=rxURT8J1XfBn3Vh3Dx_WzVADWn9zVriIty5KZEG-x6o,1622
 brainstate/nn/_collective_ops.py,sha256=BzKKUxeBPtZojHqvOzgPkz9EFaVJUwQWmqR18mF2l38,6233
-brainstate/nn/_exp_euler.py,sha256=3Q50V9TlDU-peSu0Exr9cjuO8nGZxrF-qljH5qxT6yM,3520
+brainstate/nn/_exp_euler.py,sha256=yjkfSllFxGWKEAlHo5AzBizzkFj6FEVDKmFV6E2g214,3521
 brainstate/nn/_exp_euler_test.py,sha256=clwRD8QR71k1jn6NrACMDEUcFMh0J9RTosoPnlYWUkw,1242
 brainstate/nn/_module.py,sha256=HDLPvLfB7jat2VT3gBu0MxA7vfzK7xgowemitHX8Cgo,10835
 brainstate/nn/_module_test.py,sha256=V4ZhiY_zYPvArkB2eeOTtZcgQrtlRyXKMbS1AJH4vC8,8893
@@ -74,18 +82,18 @@ brainstate/nn/metrics.py,sha256=iupHjSRTHYY-HmEPBC4tXWrZfF4zh1ek2NwSAA0gnwE,1473
 brainstate/nn/_dyn_impl/__init__.py,sha256=Oazar7h89dp1WA2Vx4Tj7gCBhxJKH4LAUEABkBEG7vU,1462
 brainstate/nn/_dyn_impl/_dynamics_neuron.py,sha256=cTbIn41EPYG0h3ICzKBXxpgB6wwA2K8k5FAcf3Pa5N8,10927
 brainstate/nn/_dyn_impl/_dynamics_neuron_test.py,sha256=Tfzrzu7udGrLJGnqItiLWe5WT0dgduvYOgzGCnaPJQg,6317
-brainstate/nn/_dyn_impl/_dynamics_synapse.py,sha256=UYRdbXsuLRS_PWuF1hGXUIZEC-g0uuVzp3VSyF1rQxk,11948
+brainstate/nn/_dyn_impl/_dynamics_synapse.py,sha256=BP-ko0FyjWZopuUhAy3Ot3wWRQlGcpumWJpKrQakqok,11869
 brainstate/nn/_dyn_impl/_dynamics_synapse_test.py,sha256=t5i-HV0ii9sUNzWTEv04o26QVtQ-mCdMJcFq2MD755A,4981
 brainstate/nn/_dyn_impl/_inputs.py,sha256=6eZKnkmrM0Gog2fpSKjSnwnQvhbFYhG4q9Vuo-GH2LI,5050
 brainstate/nn/_dyn_impl/_projection_alignpost.py,sha256=PNC1Tzx_SF2DHAHeJCufXzO_Q4qLoBpWABI45B3GRuc,876
-brainstate/nn/_dyn_impl/_rate_rnns.py,sha256=TJOD3av5gnDGZlqcrxRo4chuCixrDhUsYbmPEhU-608,14773
+brainstate/nn/_dyn_impl/_rate_rnns.py,sha256=dz_yT_6hJVhKulcjIARbGtmMzZqISws96CtBc6o5GOo,14768
 brainstate/nn/_dyn_impl/_rate_rnns_test.py,sha256=gNgtr-a4ZiU1XF9wFG1HiJ9fLosfWchVR9Zn1x39xt4,2452
 brainstate/nn/_dyn_impl/_readout.py,sha256=iYk2lKkB42OClLUlXQVr8SIqL4NzwZzVE3rlEAExGvw,4370
 brainstate/nn/_dyn_impl/_readout_test.py,sha256=R9JJPRvy3mAHSv8n1Hzjk2kBSDjBzJNbS83ystll86s,2109
 brainstate/nn/_dynamics/__init__.py,sha256=j1HSWu01wf5-KjSaNhBC9utVGDALOhUsFPrLPcPPDsM,1208
 brainstate/nn/_dynamics/_dynamics_base.py,sha256=vOOi7lQQmfVUVubm_1G2Xj3kAd4S9FFtbUDKxDx97Kg,21637
 brainstate/nn/_dynamics/_dynamics_base_test.py,sha256=gXMwENqqSvyZbMpLP0QtYndJ_h39dF5gIeiiSbMAjTk,2721
-brainstate/nn/_dynamics/_projection_base.py,sha256=FXeD8UTdoEtm1_CG8ZOtzfnD8Cx6BijsxwqhG4ewdJg,12596
+brainstate/nn/_dynamics/_projection_base.py,sha256=nSdaA1UgCormDHN7LFebhHl7CjWUuJ4b3zeJr0vq5Ms,12597
 brainstate/nn/_dynamics/_state_delay.py,sha256=nZYGmVKmQvAQu-W4YOUFH1gnr-ZS3rg_GNkRhI8rQ-I,16761
 brainstate/nn/_dynamics/_synouts.py,sha256=9TGAc-nVa50th7KKn4oKLbro-4W4rwxYvp-eu7ksAIE,4491
 brainstate/nn/_dynamics/_synouts_test.py,sha256=V_jDswRN4VvEXD-2yJO3VA1TALgX0HK6oPBQiUntOWc,2266
@@ -94,19 +102,21 @@ brainstate/nn/_elementwise/_dropout.py,sha256=jxqce-lcH5P0AotHZKhdp_Ho1n7qUuvSe6
 brainstate/nn/_elementwise/_dropout_test.py,sha256=Qn7xqZOyZMPCGF6tFjTiPId0yELOXjSsW5-hgihP3fE,4383
 brainstate/nn/_elementwise/_elementwise.py,sha256=om-KpwDTk5yFG5KBYXXHquRLV7s28_FJjk-omvyMyvQ,33342
 brainstate/nn/_elementwise/_elementwise_test.py,sha256=SZI9jB39sZ5SO1dpWGW-PhodthwN0GU9FY1nqf2fWcs,5341
-brainstate/nn/_interaction/__init__.py,sha256=BmiJxuORxiGoziApKlcMTQ-Q5TyV57p5yGKFxNanDCQ,1196
-brainstate/nn/_interaction/_connections.py,sha256=46OYPtLvJZHhJcA6ZgxYKdqnjR22xftac4PQSnoBj0w,26382
-brainstate/nn/_interaction/_connections_test.py,sha256=53GHjv_YieAzwht4F9kpYZ83a5bOaLa4WRHa7e2iZf4,9076
+brainstate/nn/_interaction/__init__.py,sha256=TTY_SeNrdx4VnUSw6vdyl02OHdS9Qs15cWBp6kjsyNQ,1289
+brainstate/nn/_interaction/_conv.py,sha256=LgWO4TeKRru07UEUga3YX6xog6WHtOvKdKtgxGnHUvw,18512
+brainstate/nn/_interaction/_conv_test.py,sha256=fHXRFYnDghFiKre63RqMwIE_gbPKdK34UPhKOz-J3qU,8695
 brainstate/nn/_interaction/_embedding.py,sha256=iK0I1ExKWFa_QzV9UDGj32Ljsmdr1g_LlAtMcusebxU,2187
+brainstate/nn/_interaction/_linear.py,sha256=bjiWGJCe81ugQQOykwjWlLW5uhe0CHWwkPA20a4n5YQ,21340
+brainstate/nn/_interaction/_linear_test.py,sha256=KlvFZA0rpyaspf4LT4K7u-RR5jCEB_q1WReqAw9sFcU,1274
 brainstate/nn/_interaction/_normalizations.py,sha256=j0sI8prfK65daQBizPFBHbs57zfGXzlnN2DxKqoN8Wk,14842
 brainstate/nn/_interaction/_normalizations_test.py,sha256=2p1Jf8nA999VYGWbvOZfKYlKk6UmL0vaEB76xkXxkXw,2438
 brainstate/nn/_interaction/_poolings.py,sha256=LpwuyeNBVCaVFW7zWc7E-vvlYqx54h46Br5XT6zd_94,47020
 brainstate/nn/_interaction/_poolings_test.py,sha256=wmd5PngZ3E9tNyF3s0xk-DoDR5yFqpTi9A6nbNoIqn4,7429
 brainstate/optim/__init__.py,sha256=7Ao0LCtDNAoxSRSXiLLKnd1_4mR2GSExizpN38il-Fo,1195
 brainstate/optim/_base.py,sha256=NbP3fzVslfnmJAOWAPD7o9TDWeRdw4CRdNfnfkMTfkU,1873
-brainstate/optim/_lr_scheduler.py,sha256=4E6u86C70QTMcIipZ22YoYLQM5R3d90hCgD9FTNJZb4,15324
+brainstate/optim/_lr_scheduler.py,sha256=Tw-aH5knnMh9eAi2LdxkZ6cAk5DbEn0tW0C738xqjFA,15325
 brainstate/optim/_lr_scheduler_test.py,sha256=_amijy9WzuvVWRC4GiuzyaC_278QG97EYZ1WtsE2IyA,1778
-brainstate/optim/_optax_optimizer.py,sha256=yctoMbFdtgaih2X6JqhbxHF575EHS-lMZ8X_Y6_USVY,4825
+brainstate/optim/_optax_optimizer.py,sha256=SuXV_xUBfhOw1_C2J5TIpy3dXDtI9VJFaSMLy8hLcXE,5312
 brainstate/optim/_optax_optimizer_test.py,sha256=DAomE8Eu3dn4gh1S3EZ_u4pW4rhcl16vWPbnDcN3Rs4,1762
 brainstate/optim/_sgd_optimizer.py,sha256=NVKYhGcw2D1ksNWUIXZcj-74LUaan8XL3EERk-EHMRI,46008
 brainstate/random/__init__.py,sha256=c5q-RC3grRIjx-HBb2IhKZpi_xzbFmUUxzRAzqfREic,1045
@@ -128,8 +138,8 @@ brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7
 brainstate/util/_struct.py,sha256=0exv0oOiSt1hmx20Y4J2-pCGtCTx13WcAlEYSBkyung,17640
 brainstate/util/_tracers.py,sha256=-sX76GJRThdSpDJBejAIzDdBbVhmH6kb-1WoDJVI7V0,2556
 brainstate/util/_visualization.py,sha256=n4ZVz10z7VBqA0cKO6vyHwEMprWJgPeEqtITzDMai2Y,1519
-brainstate-0.1.0.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.dist-info/METADATA,sha256=WmOW4lSuLuiDjOBC29ngnWvQ3p7vV2VA0BCluQpt2dE,3388
-brainstate-0.1.0.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.dist-info/RECORD,,
+brainstate-0.1.0.post20241125.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20241125.dist-info/METADATA,sha256=PV2RbaJa62UxM18IcJFfmHTSt2ClcZi6LnIS1Tbky-4,3401
+brainstate-0.1.0.post20241125.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20241125.dist-info/top_level.txt,sha256=MVkn5SZk0qis32_AGxU2Tg1xADfj2IgCNS25CQD7_ng,21
+brainstate-0.1.0.post20241125.dist-info/RECORD,,

{brainstate-0.1.0.dist-info → brainstate-0.1.0.post20241125.dist-info}/top_level.txt

@@ -1 +1,2 @@
+benchmark
 brainstate

brainstate/event/_fixed_probability.py

@@ -1,271 +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 __future__ import annotations
-
-from typing import Union, Callable, Optional
-
-import brainunit as u
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-from brainstate._state import ParamState
-from brainstate._utils import set_module_as
-from brainstate.compile import for_loop
-from brainstate.init import param
-from brainstate.nn._module import Module
-from brainstate.random import RandomState
-from brainstate.typing import ArrayLike
-from ._misc import FloatScalar, IntScalar
-
-__all__ = [
-    'FixedProb',
-]
-
-
-class FixedProb(Module):
-    """
-    The FixedProb module implements a fixed probability connection with CSR sparse data structure.
-
-    Parameters
-    ----------
-    n_pre : int
-        Number of pre-synaptic neurons.
-    n_post : int
-        Number of post-synaptic neurons.
-    prob : float
-        Probability of connection.
-    weight : float or callable or jax.Array or brainunit.Quantity
-        Maximum synaptic conductance.
-    allow_multi_conn : bool, optional
-        Whether multiple connections are allowed from a single pre-synaptic neuron.
-        Default is True, meaning that a value of ``a`` can be selected multiple times.
-    prob : float
-        Probability of connection.
-    name : str, optional
-        Name of the module.
-    """
-
-    __module__ = 'brainstate.event'
-
-    def __init__(
-        self,
-        n_pre: IntScalar,
-        n_post: IntScalar,
-        prob: FloatScalar,
-        weight: Union[Callable, ArrayLike],
-        allow_multi_conn: bool = True,
-        seed: Optional[int] = None,
-        name: Optional[str] = None,
-        grad_mode: str = 'vjp'
-    ):
-        super().__init__(name=name)
-        self.n_pre = n_pre
-        self.n_post = n_post
-        self.in_size = n_pre
-        self.out_size = n_post
-
-        self.n_conn = int(n_post * prob)
-        if self.n_conn < 1:
-            raise ValueError(
-                f"The number of connections must be at least 1. Got: int({n_post} * {prob}) = {self.n_conn}")
-
-        assert grad_mode in ['vjp', 'jvp'], f"Unsupported grad_mode: {grad_mode}"
-        self.grad_mode = grad_mode
-
-        # indices of post connected neurons
-        if allow_multi_conn:
-            self.indices = np.random.RandomState(seed).randint(0, n_post, size=(self.n_pre, self.n_conn))
-        else:
-            rng = RandomState(seed)
-            self.indices = for_loop(lambda i: rng.choice(n_post, size=(self.n_conn,), replace=False), np.arange(n_pre))
-        self.indices = u.math.asarray(self.indices)
-
-        # maximum synaptic conductance
-        weight = param(weight, (self.n_pre, self.n_conn), allow_none=False)
-        self.weight = ParamState(weight)
-
-    def update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
-        device_kind = jax.devices()[0].platform  # spk.device.device_kind
-        if device_kind == 'cpu':
-            return cpu_fixed_prob(self.indices,
-                                  u.math.asarray(self.weight.value),
-                                  u.math.asarray(spk),
-                                  n_post=self.n_post,
-                                  grad_mode=self.grad_mode)
-        elif device_kind in ['gpu', 'tpu']:
-            raise NotImplementedError()
-        else:
-            raise ValueError(f"Unsupported device: {device_kind}")
-
-
-@set_module_as('brainstate.event')
-def cpu_fixed_prob(
-    indices: jax.Array,
-    weight: Union[u.Quantity, jax.Array],
-    spk: jax.Array,
-    *,
-    n_post: int,
-    grad_mode: str = 'vjp'
-) -> Union[u.Quantity, jax.Array]:
-    """
-    The FixedProb module implements a fixed probability connection with CSR sparse data structure.
-
-    Parameters
-    ----------
-    n_post : int
-        Number of post-synaptic neurons.
-    weight : brainunit.Quantity or jax.Array
-        Maximum synaptic conductance.
-    spk : jax.Array
-        Spike events.
-    indices : jax.Array
-        Indices of post connected neurons.
-    grad_mode : str, optional
-        Gradient mode. Default is 'vjp'. Can be 'vjp' or 'jvp'.
-
-    Returns
-    -------
-    post_data : brainunit.Quantity or jax.Array
-        Post synaptic data.
-    """
-    unit = u.get_unit(weight)
-    weight = u.get_mantissa(weight)
-    indices = jnp.asarray(indices)
-    spk = jnp.asarray(spk)
-
-    def mv(spk_vector):
-        assert spk_vector.ndim == 1, f"spk must be 1D. Got: {spk.ndim}"
-        if grad_mode == 'vjp':
-            post_data = _cpu_event_fixed_prob_mv_vjp(indices, weight, spk_vector, n_post)
-        elif grad_mode == 'jvp':
-            post_data = _cpu_event_fixed_prob_mv_jvp(indices, weight, spk_vector, n_post)
-        else:
-            raise ValueError(f"Unsupported grad_mode: {grad_mode}")
-        return post_data
-
-    assert spk.ndim >= 1, f"spk must be at least 1D. Got: {spk.ndim}"
-    assert weight.ndim in [2, 0], f"weight must be 2D or 0D. Got: {weight.ndim}"
-    assert indices.ndim == 2, f"indices must be 2D. Got: {indices.ndim}"
-
-    if spk.ndim == 1:
-        post_data = mv(spk)
-    else:
-        shape = spk.shape[:-1]
-        post_data = jax.vmap(mv)(u.math.reshape(spk, (-1, spk.shape[-1])))
-        post_data = u.math.reshape(post_data, shape + post_data.shape[-1:])
-    return u.maybe_decimal(u.Quantity(post_data, unit=unit))
-
-
-# -------------------
-# CPU Implementation
-# -------------------
-
-
-def _cpu_event_fixed_prob_mv(indices, g_max, spk, n_post: int) -> jax.Array:
-    def scan_fn(post, i):
-        w = g_max if jnp.size(g_max) == 1 else g_max[i]
-        ids = indices[i]
-        sp = spk[i]
-        if spk.dtype == jnp.bool_:
-            post = jax.lax.cond(sp, lambda: post.at[ids].add(w), lambda: post)
-        else:
-            post = jax.lax.cond(sp == 0., lambda: post, lambda: post.at[ids].add(w * sp))
-        return post, None
-
-    return jax.lax.scan(scan_fn, jnp.zeros((n_post,), dtype=g_max.dtype), np.arange(len(spk)))[0]
-
-
-# --------------
-# VJP
-# --------------
-
-def _cpu_event_fixed_prob_mv_fwd(indices, g_max, spk, n_post):
-    return _cpu_event_fixed_prob_mv(indices, g_max, spk, n_post=n_post), (g_max, spk)
-
-
-def _cpu_event_fixed_prob_mv_bwd(indices, n_post, res, ct):
-    weight, spk = res
-
-    # ∂L/∂spk = ∂L/∂y * ∂y/∂spk
-    homo = jnp.size(weight) == 1
-    if homo:  # homogeneous weight
-        ct_spk = jax.vmap(lambda idx: jnp.sum(ct[idx] * weight))(indices)
-    else:  # heterogeneous weight
-        ct_spk = jax.vmap(lambda idx, w: jnp.inner(ct[idx], w))(indices, weight)
-
-    # ∂L/∂w = ∂L/∂y * ∂y/∂w
-    if homo:  # scalar
-        ct_gmax = _cpu_event_fixed_prob_mv(indices, jnp.asarray(1.), spk, n_post=n_post)
-        ct_gmax = jnp.inner(ct, ct_gmax)
-    else:
-        def scan_fn(d_gmax, i):
-            if spk.dtype == jnp.bool_:
-                d_gmax = jax.lax.cond(spk[i], lambda: d_gmax.at[i].add(ct[indices[i]]), lambda: d_gmax)
-            else:
-                d_gmax = jax.lax.cond(spk[i] == 0., lambda: d_gmax, lambda: d_gmax.at[i].add(ct[indices[i]] * spk[i]))
-            return d_gmax, None
-
-        ct_gmax = jax.lax.scan(scan_fn, jnp.zeros_like(weight), np.arange(len(spk)))[0]
-    return ct_gmax, ct_spk
-
-
-_cpu_event_fixed_prob_mv_vjp = jax.custom_vjp(_cpu_event_fixed_prob_mv, nondiff_argnums=(0, 3))
-_cpu_event_fixed_prob_mv_vjp.defvjp(_cpu_event_fixed_prob_mv_fwd, _cpu_event_fixed_prob_mv_bwd)
-
-
-# --------------
-# JVP
-# --------------
-
-
-def _cpu_event_fixed_prob_mv_jvp_rule(indices, n_post, primals, tangents):
-    # forward pass
-    weight, spk = primals
-    y = _cpu_event_fixed_prob_mv(indices, weight, spk, n_post=n_post)
-
-    # forward gradients
-    gmax_dot, spk_dot = tangents
-
-    # ∂y/∂gmax
-    dgmax = _cpu_event_fixed_prob_mv(indices, gmax_dot, spk, n_post=n_post)
-
-    def scan_fn(post, i):
-        ids = indices[i]
-        w = weight if jnp.size(weight) == 1 else weight[i]
-        post = post.at[ids].add(w * spk_dot[i])
-        return post, None
-
-    # ∂y/∂gspk
-    dspk = jax.lax.scan(scan_fn, jnp.zeros((n_post,), dtype=weight.dtype), np.arange(len(spk)))[0]
-    return y, dgmax + dspk
-
-
-_cpu_event_fixed_prob_mv_jvp = jax.custom_jvp(_cpu_event_fixed_prob_mv, nondiff_argnums=(0, 3))
-_cpu_event_fixed_prob_mv_jvp.defjvp(_cpu_event_fixed_prob_mv_jvp_rule)
-
-
-def _gpu_event_fixed_prob_mv(indices, g_max, spk, n_post: int) -> jax.Array:
-    def scan_fn(post, i):
-        w = g_max if jnp.size(g_max) == 1 else g_max[i]
-        ids = indices[i]
-        sp = spk[i]
-        if spk.dtype == jnp.bool_:
-            post = jax.lax.cond(sp, lambda: post.at[ids].add(w), lambda: post)
-        else:
-            post = jax.lax.cond(sp == 0., lambda: post, lambda: post.at[ids].add(w * sp))
-        return post, None
-
-    return jax.lax.scan(scan_fn, jnp.zeros((n_post,), dtype=g_max.dtype), np.arange(len(spk)))[0]

brainstate/event/_linear.py

@@ -1,219 +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 __future__ import annotations
-
-from typing import Union, Callable, Optional
-
-import brainunit as u
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-from brainstate._state import ParamState, State
-from brainstate._utils import set_module_as
-from brainstate.init import param
-from brainstate.nn._module import Module
-from brainstate.typing import ArrayLike
-from ._misc import IntScalar
-
-__all__ = [
-    'Linear',
-]
-
-
-class Linear(Module):
-    """
-    The FixedProb module implements a fixed probability connection with CSR sparse data structure.
-
-    Parameters
-    ----------
-    n_pre : int
-        Number of pre-synaptic neurons.
-    n_post : int
-        Number of post-synaptic neurons.
-    weight : float or callable or jax.Array or brainunit.Quantity
-        Maximum synaptic conductance.
-    name : str, optional
-        Name of the module.
-    """
-
-    __module__ = 'brainstate.event'
-
-    def __init__(
-        self,
-        n_pre: IntScalar,
-        n_post: IntScalar,
-        weight: Union[Callable, ArrayLike],
-        name: Optional[str] = None,
-        grad_mode: str = 'vjp'
-    ):
-        super().__init__(name=name)
-        self.n_pre = n_pre
-        self.n_post = n_post
-        self.in_size = n_pre
-        self.out_size = n_post
-
-        assert grad_mode in ['vjp', 'jvp'], f"Unsupported grad_mode: {grad_mode}"
-        self.grad_mode = grad_mode
-
-        # maximum synaptic conductance
-        weight = param(weight, (self.n_pre, self.n_post), allow_none=False)
-        self.weight = ParamState(weight)
-
-    def update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
-        weight = self.weight.value if isinstance(self.weight, State) else self.weight
-        if u.math.size(weight) == 1:
-            return u.math.ones(self.n_post) * (u.math.sum(spk) * weight)
-
-        device_kind = jax.devices()[0].platform  # spk.device.device_kind
-        if device_kind == 'cpu':
-            return cpu_event_linear(u.math.asarray(weight),
-                                    u.math.asarray(spk),
-                                    n_post=self.n_post,
-                                    grad_mode=self.grad_mode)
-        elif device_kind in ['gpu', 'tpu']:
-            raise NotImplementedError()
-        else:
-            raise ValueError(f"Unsupported device: {device_kind}")
-
-
-@set_module_as('brainstate.event')
-def cpu_event_linear(
-    g_max: Union[u.Quantity, jax.Array],
-    spk: jax.Array,
-    *,
-    n_post: int = None,
-    grad_mode: str = 'vjp'
-) -> Union[u.Quantity, jax.Array]:
-    """
-    The FixedProb module implements a fixed probability connection with CSR sparse data structure.
-
-    Parameters
-    ----------
-    n_post : int
-        Number of post-synaptic neurons.
-    g_max : brainunit.Quantity or jax.Array
-        Maximum synaptic conductance.
-    spk : jax.Array
-        Spike events.
-    grad_mode : str, optional
-        Gradient mode. Default is 'vjp'. Can be 'vjp' or 'jvp'.
-
-    Returns
-    -------
-    post_data : brainunit.Quantity or jax.Array
-        Post synaptic data.
-    """
-    unit = u.get_unit(g_max)
-    g_max = u.get_mantissa(g_max)
-    spk = jnp.asarray(spk)
-
-    def mv(spk_vector):
-        assert spk_vector.ndim == 1, f"spk must be 1D. Got: {spk.ndim}"
-        if jnp.size(g_max) == 1:
-            assert isinstance(n_post, int), f"n_post must be an integer when weight is homogenous. Got: {n_post}"
-            # return jnp.full((n_post,), fill_value=jnp.sum(spk_vector) * weight)
-            return jnp.ones((n_post,), dtype=g_max.dtype) * (jnp.sum(spk_vector) * g_max)
-
-        if grad_mode == 'vjp':
-            post = _cpu_event_linear_mv_vjp(g_max, spk_vector)
-        elif grad_mode == 'jvp':
-            post = _cpu_event_linear_mv_jvp(g_max, spk_vector)
-        else:
-            raise ValueError(f"Unsupported grad_mode: {grad_mode}")
-        return post
-
-    assert spk.ndim >= 1, f"spk must be at least 1D. Got: {spk.ndim}"
-    assert g_max.ndim in [2, 0], f"weight must be 2D or 0D. Got: {g_max.ndim}"
-
-    if spk.ndim == 1:
-        post_data = mv(spk)
-    else:
-        shape = spk.shape[:-1]
-        post_data = jax.vmap(mv)(u.math.reshape(spk, (-1, spk.shape[-1])))
-        post_data = u.math.reshape(post_data, shape + post_data.shape[-1:])
-    return u.maybe_decimal(u.Quantity(post_data, unit=unit))
-
-
-# --------------
-# Implementation
-# --------------
-
-
-def _cpu_event_linear_mv(g_max, spk) -> jax.Array:
-    def scan_fn(post, i):
-        sp = spk[i]
-        if spk.dtype == jnp.bool_:
-            post = jax.lax.cond(sp, lambda: post + g_max[i], lambda: post)
-        else:
-            post = jax.lax.cond(sp == 0., lambda: post, lambda: post + g_max[i] * sp)
-        return post, None
-
-    return jax.lax.scan(scan_fn, jnp.zeros(g_max.shape[1], dtype=g_max.dtype), np.arange(len(spk)))[0]
-
-
-# --------------
-# VJP
-# --------------
-
-def _cpu_event_linear_mv_fwd(g_max, spk):
-    return _cpu_event_linear_mv(g_max, spk), (g_max, spk)
-
-
-def _cpu_event_linear_mv_bwd(res, ct):
-    g_max, spk = res
-
-    # ∂L/∂spk = ∂L/∂y * ∂y/∂spk
-    ct_spk = jnp.matmul(g_max, ct)
-
-    # ∂L/∂w = ∂L/∂y * ∂y/∂w
-    def map_fn(sp):
-        if spk.dtype == jnp.bool_:
-            d_gmax = jax.lax.cond(sp, lambda: ct, lambda: jnp.zeros_like(ct))
-        else:
-            d_gmax = jax.lax.cond(sp == 0., lambda: jnp.zeros_like(ct), lambda: ct * sp)
-        return d_gmax
-
-    ct_gmax = jax.vmap(map_fn)(spk)
-    return ct_gmax, ct_spk
-
-
-_cpu_event_linear_mv_vjp = jax.custom_vjp(_cpu_event_linear_mv)
-_cpu_event_linear_mv_vjp.defvjp(_cpu_event_linear_mv_fwd, _cpu_event_linear_mv_bwd)
-
-
-# --------------
-# JVP
-# --------------
-
-
-def _cpu_event_linear_mv_jvp_rule(primals, tangents):
-    # forward pass
-    g_max, spk = primals
-    y = _cpu_event_linear_mv(g_max, spk)
-
-    # forward gradients
-    gmax_dot, spk_dot = tangents
-
-    # ∂y/∂gmax
-    dgmax = _cpu_event_linear_mv(gmax_dot, spk)
-
-    # ∂y/∂gspk
-    dspk = spk_dot @ g_max
-    return y, dgmax + dspk
-
-
-_cpu_event_linear_mv_jvp = jax.custom_jvp(_cpu_event_linear_mv)
-_cpu_event_linear_mv_jvp.defjvp(_cpu_event_linear_mv_jvp_rule)