brainstate 0.0.2.post20240913__py2.py3-none-any.whl → 0.0.2.post20241010__py2.py3-none-any.whl

brainstate/nn/event/fixed_probability.py

@@ -0,0 +1,276 @@
+# 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.random import RandomState
+from brainstate.transform import for_loop
+from brainstate.typing import ArrayLike
+from ._misc import FloatScalar, IntScalar
+
+__all__ = [
+  'EventFixedProb',
+]
+
+
+class EventFixedProb(DnnLayer):
+  """
+  The EventFixedProb 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.
+  mode : brainstate.mixin.Mode, optional
+      Mode of the module.
+  """
+
+  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,
+      mode: Optional[Mode] = None,
+      grad_mode: str = 'vjp'
+  ):
+    super().__init__(name=name, mode=mode)
+    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))
+
+    # maximum synaptic conductance
+    weight = param(weight, (self.n_pre, self.n_conn), allow_none=False)
+    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
+    device_kind = jax.devices()[0].platform  # spk.device.device_kind
+    if device_kind == 'cpu':
+      return cpu_event_fixed_prob(u.math.asarray(self.indices),
+                                  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}")
+
+
+def cpu_event_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 EventFixedProb 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/nn/event/fixed_probability_test.py

@@ -0,0 +1,127 @@
+# 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
+from absl.testing import parameterized
+
+import brainstate as bst
+from brainstate.nn.event.fixed_probability import EventFixedProb
+
+
+class TestFixedProbCSR(parameterized.TestCase):
+  @parameterized.product(
+    allow_multi_conn=[True, False]
+  )
+  def test1(self, allow_multi_conn):
+    x = bst.random.rand(20) < 0.1
+    # x = bst.random.rand(20)
+    m = EventFixedProb(20, 40, 0.1, 1.0, seed=123, allow_multi_conn=allow_multi_conn)
+    y = m(x)
+    print(y)
+
+    m2 = EventFixedProb(20, 40, 0.1, bst.init.KaimingUniform(), seed=123)
+    print(m2(x))
+
+  def test_grad_bool(self):
+    n_in = 20
+    n_out = 30
+    x = bst.random.rand(n_in) < 0.3
+    fn = EventFixedProb(n_in, n_out, 0.1, bst.init.KaimingUniform(), seed=123)
+
+    def f(x):
+      return fn(x).sum()
+
+    with self.assertRaises(TypeError):
+      print(jax.grad(f)(x))
+
+  @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)
+
+    if homo_w:
+      fn = EventFixedProb(n_in, n_out, 0.1, 1.5, seed=123)
+    else:
+      fn = EventFixedProb(n_in, n_out, 0.1, bst.init.KaimingUniform(), seed=123)
+    w = fn.weight
+
+    def f(x, w):
+      fn.weight = w
+      return fn(x).sum()
+
+    r = bst.transform.grad(f, argnums=(0, 1))(x, w)
+
+    # -------------------
+    # TRUE gradients
+
+    def true_fn(x, w, indices, n_post):
+      post = jnp.zeros((n_post,))
+      for i in range(n_in):
+        post = post.at[indices[i]].add(w * x[i] if homo_w else w[i] * x[i])
+      return post
+
+    def f2(x, w):
+      return true_fn(x, w, fn.indices, 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]))
+    print(r[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)
+
+    fn = EventFixedProb(n_in, n_out, 0.1, 1.5 if homo_w else bst.init.KaimingUniform(), seed=123, 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 true_fn(x, w, indices, n_post):
+      post = jnp.zeros((n_post,))
+      for i in range(n_in):
+        post = post.at[indices[i]].add(w * x[i] if homo_w else w[i] * x[i])
+      return post
+
+    def f2(x, w):
+      return true_fn(x, w, fn.indices, 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))

brainstate/nn/event/linear.py

@@ -0,0 +1,220 @@
+# 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__ = [
+  'EventDense',
+]
+
+
+class EventDense(DnnLayer):
+  """
+  The EventFixedProb 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,
+      weight: Union[Callable, ArrayLike],
+      name: Optional[str] = None,
+      mode: Optional[Mode] = None,
+      grad_mode: str = 'vjp'
+  ):
+    super().__init__(name=name, mode=mode)
+    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)
+    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 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}")
+
+
+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 EventFixedProb 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)