brainstate 0.0.2.post20241010__py2.py3-none-any.whl → 0.1.0.post20241122__py2.py3-none-any.whl

brainstate/nn/event/fixed_probability_test.py

@@ -1,127 +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.
-# ==============================================================================
-
-
-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

@@ -1,220 +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
-
-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)

brainstate/nn/event/linear_test.py

@@ -1,111 +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.
-# ==============================================================================
-
-
-import jax
-import jax.numpy as jnp
-from absl.testing import parameterized
-
-import brainstate as bst
-from brainstate.nn.event.linear import EventDense
-
-
-class TestEventLinear(parameterized.TestCase):
-  @parameterized.product(
-    homo_w=[True, False],
-    bool_x=[True, False],
-  )
-  def test1(self, homo_w, bool_x):
-    x = bst.random.rand(20) < 0.1
-    if not bool_x:
-      x = jnp.asarray(x, dtype=float)
-    m = EventDense(20, 40, 1.5 if homo_w else bst.init.KaimingUniform())
-    y = m(x)
-    print(y)
-
-    self.assertTrue(jnp.allclose(y, (x.sum() * m.weight) if homo_w else (x @ m.weight)))
-
-  def test_grad_bool(self):
-    n_in = 20
-    n_out = 30
-    x = bst.random.rand(n_in) < 0.3
-    fn = EventDense(n_in, n_out, bst.init.KaimingUniform())
-
-    with self.assertRaises(TypeError):
-      print(jax.grad(lambda x: fn(x).sum())(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)
-
-    fn = EventDense(n_in, n_out, 1.5 if homo_w else bst.init.KaimingUniform())
-    w = fn.weight
-
-    def f(x, w):
-      fn.weight = w
-      return fn(x).sum()
-
-    r1 = jax.grad(f, argnums=(0, 1))(x, w)
-
-    # -------------------
-    # TRUE gradients
-
-    def f2(x, w):
-      y = (x @ (jnp.ones([n_in, n_out]) * w)) if homo_w else (x @ w)
-      return y.sum()
-
-    r2 = jax.grad(f2, argnums=(0, 1))(x, w)
-    self.assertTrue(jnp.allclose(r1[0], r2[0]))
-    self.assertTrue(jnp.allclose(r1[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)
-
-    fn = EventDense(n_in, n_out, 1.5 if homo_w else bst.init.KaimingUniform(), 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):
-      y = (x @ (jnp.ones([n_in, n_out]) * w)) if homo_w else (x @ w)
-      return y
-
-    o2, r2 = jax.jvp(f, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
-    self.assertTrue(jnp.allclose(o1, o2))
-    self.assertTrue(jnp.allclose(r1, r2))