brainstate 0.0.2.post20240910__py2.py3-none-any.whl → 0.0.2.post20241009__py2.py3-none-any.whl

brainstate/nn/_poolings.py

@@ -21,7 +21,7 @@ import functools
 from typing import Sequence, Optional
 from typing import Union, Tuple, Callable, List
 
-import brainunit as bu
+import brainunit as u
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -29,7 +29,7 @@ import numpy as np
 from ._base import DnnLayer, ExplicitInOutSize
 from .. import environ
 from ..mixin import Mode
-from ..typing import Size
+from brainstate.typing import Size
 
 __all__ = [
   'Flatten', 'Unflatten',
@@ -85,7 +85,7 @@ class Flatten(DnnLayer, ExplicitInOutSize):
 
     if in_size is not None:
       self.in_size = tuple(in_size)
-      y = jax.eval_shape(functools.partial(bu.math.flatten, start_axis=start_axis, end_axis=end_axis),
+      y = jax.eval_shape(functools.partial(u.math.flatten, start_axis=start_axis, end_axis=end_axis),
                          jax.ShapeDtypeStruct(self.in_size, environ.dftype()))
       self.out_size = y.shape
 
@@ -101,7 +101,7 @@ class Flatten(DnnLayer, ExplicitInOutSize):
         start_axis = self.start_axis + dim_diff
       else:
         start_axis = x.ndim + self.start_axis
-    return bu.math.flatten(x, start_axis, self.end_axis)
+    return u.math.flatten(x, start_axis, self.end_axis)
 
   def __repr__(self) -> str:
     return f'{self.__class__.__name__}(start_axis={self.start_axis}, end_axis={self.end_axis})'
@@ -153,12 +153,12 @@ class Unflatten(DnnLayer, ExplicitInOutSize):
 
     if in_size is not None:
       self.in_size = tuple(in_size)
-      y = jax.eval_shape(functools.partial(bu.math.unflatten, axis=axis, sizes=sizes),
+      y = jax.eval_shape(functools.partial(u.math.unflatten, axis=axis, sizes=sizes),
                          jax.ShapeDtypeStruct(self.in_size, environ.dftype()))
       self.out_size = y.shape
 
   def update(self, x):
-    return bu.math.unflatten(x, self.axis, self.sizes)
+    return u.math.unflatten(x, self.axis, self.sizes)
 
   def __repr__(self):
     return f'{self.__class__.__name__}(axis={self.axis}, sizes={self.sizes})'

brainstate/nn/_rate_rnns.py

@@ -27,7 +27,7 @@ from .. import random, init, functional
 from .._module import Module
 from .._state import ShortTermState, ParamState
 from ..mixin import DelayedInit, Mode
-from ..typing import ArrayLike
+from brainstate.typing import ArrayLike
 
 __all__ = [
   'RNNCell', 'ValinaRNNCell', 'GRUCell', 'MGUCell', 'LSTMCell', 'URLSTMCell',

brainstate/nn/_readout.py

@@ -29,7 +29,7 @@ from ._misc import exp_euler_step
 from .. import environ, init, surrogate
 from .._state import ShortTermState, ParamState
 from ..mixin import Mode
-from ..typing import Size, ArrayLike, DTypeLike
+from brainstate.typing import Size, ArrayLike, DTypeLike
 
 __all__ = [
   'LeakyRateReadout',

brainstate/nn/_synouts.py

@@ -23,7 +23,7 @@ import jax.numpy as jnp
 
 from .._module import Module
 from ..mixin import DelayedInit, BindCondData
-from ..typing import ArrayLike
+from brainstate.typing import ArrayLike
 
 __all__ = [
   'SynOut', 'COBA', 'CUBA', 'MgBlock',

brainstate/nn/event/__init__.py

@@ -0,0 +1,25 @@
+# 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 .csr import *
+from .csr import __all__ as __all_csr
+from .fixed_probability import *
+from .fixed_probability import __all__ as __all_fixed_probability
+from .linear import *
+from .linear import __all__ as __all_linear
+
+__all__ = __all_fixed_probability + __all_linear + __all_csr
+del __all_fixed_probability, __all_linear, __all_csr

brainstate/nn/event/_misc.py

@@ -0,0 +1,34 @@
+# 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
+
+import numpy as np
+
+__all__ = [
+  'FloatScalar',
+  'IntScalar',
+]
+
+FloatScalar = Union[
+  np.number,  # NumPy scalar types
+  float,  # Python scalar types
+]
+
+IntScalar = Union[
+  np.number,  # NumPy scalar types
+  int,  # Python scalar types
+]

brainstate/nn/event/csr.py

@@ -0,0 +1,312 @@
+# 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
+
+import brainunit as u
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from brainstate._state import ParamState, State
+from brainstate.init import param
+from brainstate.mixin import Mode, Training
+from brainstate.nn._base import DnnLayer
+from brainstate.typing import ArrayLike
+from ._misc import IntScalar
+
+__all__ = [
+  'EventCSR',
+]
+
+
+class EventCSR(DnnLayer):
+  """
+  The EventCSR 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.
+  mode : brainstate.mixin.Mode, optional
+      Mode of the module.
+  """
+
+  def __init__(
+      self,
+      n_pre: IntScalar,
+      n_post: IntScalar,
+      indptr: ArrayLike,
+      indices: ArrayLike,
+      weight: Union[Callable, ArrayLike],
+      name: Optional[str] = None,
+      mode: Optional[Mode] = None,
+      grad_mode: str = 'vjp'
+  ):
+    super().__init__(name=name, mode=mode)
+
+    self.in_size = n_pre
+    self.out_size = n_post
+    self.n_pre = n_pre
+    self.n_post = n_post
+
+    # gradient mode
+    assert grad_mode in ['vjp', 'jvp'], f"Unsupported grad_mode: {grad_mode}"
+    self.grad_mode = grad_mode
+
+    # CSR data structure
+    indptr = jnp.asarray(indptr)
+    indices = jnp.asarray(indices)
+    assert indptr.ndim == 1, f"indptr must be 1D. Got: {indptr.ndim}"
+    assert indices.ndim == 1, f"indices must be 1D. Got: {indices.ndim}"
+    assert indptr.size == n_pre + 1, f"indptr must have size {n_pre + 1}. Got: {indptr.size}"
+    self.indptr = indptr
+    self.indices = indices
+
+    # maximum synaptic conductance
+    weight = param(weight, (len(indices),), allow_none=False)
+    # if callable(weight):
+    #   pass
+    # else:
+    #   if u.math.size(weight) != 1 and u.math.size(weight) != len(self.indices):
+    #     raise ValueError(f"weight must be 1D or 2D with size {len(self.indices)}. Got: {u.math.size(weight)}")
+    if self.mode.has(Training):
+      weight = ParamState(weight)
+    self.weight = 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 len(self.indices) == 0:
+      r = u.math.zeros(spk.shape[:-1] + (self.n_post,),
+                       dtype=weight.dtype,
+                       unit=u.get_unit(weight) * u.get_unit(spk))
+      return u.maybe_decimal(r)
+
+    device_kind = jax.devices()[0].platform  # spk.device.device_kind
+    if device_kind == 'cpu':
+      return cpu_event_csr(
+        u.math.asarray(spk),
+        u.math.asarray(self.indptr),
+        u.math.asarray(self.indices),
+        u.math.asarray(weight),
+        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}")
+
+
+def cpu_event_csr(
+    spk: jax.Array,
+    indptr: jax.Array,
+    indices: jax.Array,
+    weight: Union[u.Quantity, jax.Array],
+    *,
+    n_post: int,
+    grad_mode: str = 'vjp'
+) -> Union[u.Quantity, jax.Array]:
+  """
+  The EventCSR module implements a fixed probability connection with CSR sparse data structure.
+
+  Parameters
+  ----------
+  spk : jax.Array
+      Spike events.
+  indptr : jax.Array
+      Index pointer of post connected neurons.
+  indices : jax.Array
+      Indices of post connected neurons.
+  weight : brainunit.Quantity or jax.Array
+      Maximum synaptic conductance.
+  n_post : int
+      Number of post-synaptic 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)
+
+  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_csr_mv_vjp(spk_vector, indptr, indices, weight, n_post)
+    elif grad_mode == 'jvp':
+      post_data = _cpu_event_csr_mv_jvp(spk_vector, indptr, indices, weight, 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 [1, 0], f"g_max must be 1D or 0D. Got: {weight.ndim}"
+  assert indices.ndim == 1, f"indices must be 1D. 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))
+
+
+# --------------
+# Implementation
+# --------------
+
+
+def _cpu_event_csr_mv(
+    spk: jax.Array,
+    indptr: jax.Array,
+    indices: jax.Array,
+    weight: Union[u.Quantity, jax.Array],
+    n_post: int
+) -> jax.Array:
+  bool_x = spk.dtype == jnp.bool_
+  homo_w = jnp.size(weight) == 1
+
+  def add_fn(post_val, i_start, i_end, sp):
+    def body_fn(x):
+      post, i = x
+      i_post = indices[i]
+      w = weight if homo_w else weight[i]
+      w = w if bool_x else w * sp
+      post = post.at[i_post].add(w)
+      return post, i + 1
+
+    return jax.lax.while_loop(lambda x: x[1] < i_end, body_fn, (post_val, i_start))[0]
+
+  def scan_fn(post, i):
+    sp = spk[i]  # pre-synaptic spike event
+    if bool_x:
+      post = jax.lax.cond(sp, lambda: add_fn(post, indptr[i], indptr[i + 1], sp), lambda: post)
+    else:
+      post = jax.lax.cond(sp == 0., lambda: post, lambda: add_fn(post, indptr[i], indptr[i + 1], sp))
+    return post, None
+
+  return jax.lax.scan(scan_fn, jnp.zeros((n_post,), dtype=weight.dtype), np.arange(len(spk)))[0]
+
+
+# --------------
+# VJP
+# --------------
+
+def _cpu_event_csr_mv_fwd(
+    spk: jax.Array,
+    indptr: jax.Array,
+    indices: jax.Array,
+    weight: Union[u.Quantity, jax.Array],
+    n_post: int
+):
+  return _cpu_event_csr_mv(spk, indptr, indices, weight, n_post=n_post), (spk, weight)
+
+
+def _cpu_event_csr_mv_bwd(indptr, indices, n_post, res, ct):
+  spk, weight = res
+  homo = jnp.size(weight) == 1
+  bool_spk = spk.dtype == jnp.bool_
+
+  # ∂L/∂spk = ∂L/∂y * ∂y/∂spk
+  def fn_spk(i_pre):
+    def body_fn(x):
+      r, i = x
+      i_post = indices[i]
+      r = r + (ct[i_post] if homo else ct[i_post] * weight[i])
+      return r, i + 1
+
+    p = jax.lax.while_loop(lambda x: x[1] < indptr[i_pre + 1], body_fn, (0., indptr[i_pre]))[0]
+    p = p * weight if homo else p
+    return p
+
+  ct_spk = jax.vmap(fn_spk)(np.arange(len(spk)))
+
+  # ∂L/∂w = ∂L/∂y * ∂y/∂w
+  if homo:  # scalar
+    ct_gmax = _cpu_event_csr_mv(spk, indptr, indices, jnp.asarray(1.), n_post=n_post)
+    ct_gmax = jnp.inner(ct, ct_gmax)
+  else:
+    def single_post(dw, i_pre):
+      def body_fn(x):
+        dw, i = x
+        i_post = indices[i]
+        dw = dw.at[i].add(ct[i_post] if bool_spk else ct[i_post] * spk[i_pre])
+        return dw, i + 1
+
+      return jax.lax.while_loop(lambda x: x[1] < indptr[i_pre + 1], body_fn, (dw, indptr[i_pre]))[0]
+
+    def fn_w(dw, i_pre):
+      sp = spk[i_pre]
+      if bool_spk:
+        return jax.lax.cond(sp, lambda: single_post(dw, i_pre), lambda: dw), None
+      else:
+        return jax.lax.cond(sp == 0., lambda: dw, lambda: single_post(dw, i_pre)), None
+
+    ct_gmax = jax.lax.scan(fn_w, jnp.zeros_like(weight), np.arange(len(spk)))[0]
+  return ct_spk, ct_gmax
+
+
+_cpu_event_csr_mv_vjp = jax.custom_vjp(_cpu_event_csr_mv, nondiff_argnums=(1, 2, 4))
+_cpu_event_csr_mv_vjp.defvjp(_cpu_event_csr_mv_fwd, _cpu_event_csr_mv_bwd)
+
+
+# --------------
+# JVP
+# --------------
+
+
+def _cpu_event_csr_mv_jvp_rule(indptr, indices, n_post, primals, tangents):
+  # forward pass
+  spk, weight = primals
+  y = _cpu_event_csr_mv(spk, indptr, indices, weight, n_post=n_post)
+
+  # forward gradients
+  spk_dot, weight_dot = tangents
+  homo_w = jnp.size(weight) == 1
+
+  # ∂y/∂gmax
+  dweight = _cpu_event_csr_mv(spk, indptr, indices, weight_dot, n_post=n_post)
+
+  # ∂y/∂gspk
+  def scan_fn(post, i_pre):
+    def while_fn(x):
+      p, i, sp = x
+      i_post = indices[i]
+      p = p.at[i_post].add(sp if homo_w else sp * weight[i])
+      return p, i + 1, sp
+
+    post = jax.lax.while_loop(lambda x: x[1] < indptr[i_pre + 1],
+                              while_fn,
+                              (post, indptr[i_pre], spk_dot[i_pre]))[0]
+
+    return post, None
+
+  dspk = jax.lax.scan(scan_fn, jnp.zeros((n_post,), dtype=weight.dtype), np.arange(len(spk)))[0]
+  dspk = (dspk * weight) if homo_w else dspk
+  return y, dweight + dspk
+
+
+_cpu_event_csr_mv_jvp = jax.custom_jvp(_cpu_event_csr_mv, nondiff_argnums=(1, 2, 4))
+_cpu_event_csr_mv_jvp.defjvp(_cpu_event_csr_mv_jvp_rule)

brainstate/nn/event/csr_test.py

@@ -0,0 +1,118 @@
+# 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.
+# ==============================================================================
+
+
+import jax.numpy
+import jax.numpy as jnp
+import numpy as np
+from absl.testing import parameterized
+
+import brainstate as bst
+from brainstate.nn.event.csr import EventCSR
+
+
+def _get_csr(n_pre, n_post, prob):
+  n_conn = int(n_post * prob)
+  indptr = np.arange(n_pre + 1) * n_conn
+  indices = np.random.randint(0, n_post, (n_pre * n_conn,))
+  return indptr, indices
+
+
+def true_fn(x, w, indices, indptr, n_out):
+  homo_w = jnp.size(w) == 1
+
+  post = jnp.zeros((n_out,))
+  for i_pre in range(x.shape[0]):
+    ids = indices[indptr[i_pre]: indptr[i_pre + 1]]
+    post = post.at[ids].add(w * x[i_pre] if homo_w else w[indptr[i_pre]: indptr[i_pre + 1]] * x[i_pre])
+  return post
+
+
+class TestFixedProbCSR(parameterized.TestCase):
+  @parameterized.product(
+    homo_w=[True, False],
+  )
+  def test1(self, homo_w):
+    x = bst.random.rand(20) < 0.1
+    indptr, indices = _get_csr(20, 40, 0.1)
+    m = EventCSR(20, 40, indptr, indices, 1.5 if homo_w else bst.init.Normal())
+    y = m(x)
+    y2 = true_fn(x, m.weight, indices, indptr, 40)
+    self.assertTrue(jnp.allclose(y, y2))
+
+  @parameterized.product(
+    bool_x=[True, False],
+    homo_w=[True, False]
+  )
+  def test_vjp(self, bool_x, homo_w):
+    n_in = 20
+    n_out = 30
+    if bool_x:
+      x = jax.numpy.asarray(bst.random.rand(n_in) < 0.3, dtype=float)
+    else:
+      x = bst.random.rand(n_in)
+
+    indptr, indices = _get_csr(n_in, n_out, 0.1)
+    fn = bst.nn.EventCSR(n_in, n_out, indptr, indices, 1.5 if homo_w else bst.init.Normal())
+    w = fn.weight
+
+    def f(x, w):
+      fn.weight = w
+      return fn(x).sum()
+
+    r = jax.grad(f, argnums=(0, 1))(x, w)
+
+    # -------------------
+    # TRUE gradients
+
+    def f2(x, w):
+      return true_fn(x, w, indices, indptr, n_out).sum()
+
+    r2 = jax.grad(f2, argnums=(0, 1))(x, w)
+    self.assertTrue(jnp.allclose(r[0], r2[0]))
+    self.assertTrue(jnp.allclose(r[1], r2[1]))
+
+  @parameterized.product(
+    bool_x=[True, False],
+    homo_w=[True, False]
+  )
+  def test_jvp(self, bool_x, homo_w):
+    n_in = 20
+    n_out = 30
+    if bool_x:
+      x = jax.numpy.asarray(bst.random.rand(n_in) < 0.3, dtype=float)
+    else:
+      x = bst.random.rand(n_in)
+
+    indptr, indices = _get_csr(n_in, n_out, 0.1)
+    fn = EventCSR(n_in, n_out, indptr, indices,
+                  1.5 if homo_w else bst.init.Normal(), grad_mode='jvp')
+    w = fn.weight
+
+    def f(x, w):
+      fn.weight = w
+      return fn(x)
+
+    o1, r1 = jax.jvp(f, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
+
+    # -------------------
+    # TRUE gradients
+
+    def f2(x, w):
+      return true_fn(x, w, indices, indptr, n_out)
+
+    o2, r2 = jax.jvp(f2, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
+    self.assertTrue(jnp.allclose(r1, r2))
+    self.assertTrue(jnp.allclose(o1, o2))