brainstate 0.1.0.post20241210__py2.py3-none-any.whl → 0.1.0.post20241220__py2.py3-none-any.whl

brainstate/nn/_dyn_impl/_inputs.py

@@ -17,17 +17,23 @@ from __future__ import annotations
 from typing import Union, Optional, Sequence, Callable
 
 import brainunit as u
+import jax
+import numpy as np
 
 from brainstate import environ, init, random
 from brainstate._state import ShortTermState
-from brainstate.compile import while_loop
-from brainstate.nn._dynamics._dynamics_base import Dynamics
+from brainstate._state import State
+from brainstate.compile import while_loop, cond
+from brainstate.nn._dynamics._dynamics_base import Dynamics, Prefetch
+from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike, Size, DTypeLike
 
 __all__ = [
     'SpikeTime',
     'PoissonSpike',
     'PoissonEncoder',
+    'PoissonInput',
+    'poisson_input',
 ]
 
 
@@ -152,3 +158,122 @@ class PoissonEncoder(Dynamics):
         spikes = random.rand(*self.varshape) <= (freqs * environ.get_dt())
         spikes = u.math.asarray(spikes, dtype=self.spk_type)
         return spikes
+
+
+class PoissonInput(Module):
+    """
+    Poisson Input to the given :py:class:`brainstate.State`.
+
+    Adds independent Poisson input to a target variable. For large
+    numbers of inputs, this is much more efficient than creating a
+    `PoissonGroup`. The synaptic events are generated randomly during the
+    simulation and are not preloaded and stored in memory. All the inputs must
+    target the same variable, have the same frequency and same synaptic weight.
+    All neurons in the target variable receive independent realizations of
+    Poisson spike trains.
+
+    Args:
+      target: The variable that is targeted by this input. Should be an instance of :py:class:`~.Variable`.
+      num_input: The number of inputs.
+      freq: The frequency of each of the inputs. Must be a scalar.
+      weight: The synaptic weight. Must be a scalar.
+      name: The target name.
+    """
+
+    def __init__(
+        self,
+        target: Prefetch,
+        indices: Union[np.ndarray, jax.Array],
+        num_input: int,
+        freq: Union[int, float],
+        weight: Union[int, float],
+        name: Optional[str] = None,
+    ):
+        super().__init__(name=name)
+
+        self.target = target
+        self.indices = indices
+        self.num_input = num_input
+        self.freq = freq
+        self.weight = weight
+
+    def update(self):
+        p = self.freq * environ.get_dt()
+        a = self.num_input * p
+        b = self.num_input * (1 - p)
+
+        target = self.target()
+        target_state = getattr(self.target.module, self.target.item)
+
+        # generate Poisson input
+        inp = cond(
+            u.math.logical_and(a > 5, b > 5),
+            lambda: random.normal(a, b * p, self.indices.shape),
+            lambda: random.binomial(self.num_input, p, self.indices.shape).astype(float)
+        )
+
+        # update target variable
+        target_state.value = target.at[self.indices].add(inp * self.weight)
+
+
+def poisson_input(
+    freq: ArrayLike,
+    num_input: int,
+    weight: ArrayLike,
+    target: State,
+    indices: Optional[Union[np.ndarray, jax.Array]] = None,
+):
+    """
+    Poisson Input to the given :py:class:`brainstate.State`.
+    """
+    assert isinstance(target, State), 'The target must be a State.'
+    p = freq * environ.get_dt()
+    a = num_input * p
+    b = num_input * (1 - p)
+    tar_val = target.value
+    if indices is None:
+        # generate Poisson input
+        inp = cond(
+            u.math.logical_and(a > 5, b > 5),
+            lambda: jax.tree.map(
+                lambda tar: random.normal(a, b * p, tar.shape),
+                tar_val,
+                is_leaf=u.math.is_quantity
+            ),
+            lambda: jax.tree.map(
+                lambda tar: random.binomial(num_input, p, tar.shape).astype(float),
+                tar_val,
+                is_leaf=u.math.is_quantity
+            )
+        )
+
+        # update target variable
+        target.value = jax.tree.map(
+            lambda x: x * weight,
+            inp,
+            is_leaf=u.math.is_quantity
+        )
+
+    else:
+        # generate Poisson input
+        inp = cond(
+            u.math.logical_and(a > 5, b > 5),
+            lambda: jax.tree.map(
+                lambda tar: random.normal(a, b * p, tar[indices].shape),
+                tar_val,
+                is_leaf=u.math.is_quantity
+            ),
+            lambda: jax.tree.map(
+                lambda tar: random.binomial(num_input, p, tar[indices].shape).astype(float),
+                tar_val,
+                is_leaf=u.math.is_quantity
+            )
+        )
+
+        # update target variable
+        target.value = jax.tree.map(
+            lambda x, tar: tar.at[indices].add(x * weight),
+            inp,
+            tar_val,
+            is_leaf=u.math.is_quantity
+        )

brainstate/nn/_dynamics/_dynamics_base.py

@@ -107,6 +107,8 @@ class Dynamics(Module):
 
     __module__ = 'brainstate.nn'
 
+    graph_invisible_attrs = ('_before_updates', '_after_updates', '_current_inputs', '_delta_inputs')
+
     # before updates
     _before_updates: Optional[Dict[Hashable, Callable]]
 
@@ -443,6 +445,16 @@ class Prefetch(Node):
         item = _get_prefetch_item(self)
         return item.value if isinstance(item, State) else item
 
+    def get_item_value(self):
+        item = _get_prefetch_item(self)
+        return item.value if isinstance(item, State) else item
+
+    def get_item(self):
+        """
+        Get
+        """
+        return _get_prefetch_item(self)
+
 
 class PrefetchDelay(Node):
     def __init__(self, module: Dynamics, item: str):

brainstate/nn/_dynamics/_projection_base.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 from __future__ import annotations
 
-from typing import Union, Callable
+from typing import Union, Callable, Optional
 
 from brainstate._state import State
 from brainstate.mixin import AlignPost, ParamDescriber, BindCondData, JointTypes
@@ -60,24 +60,28 @@ def is_instance(x, cls) -> bool:
     return isinstance(x, cls)
 
 
-def get_post_repr(syn, out):
-    return f'{syn.identifier} // {out.identifier}'
+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
+    proj_name: str,
+    label: str,
 ):
     # synapse and output initialization
-    _post_repr = get_post_repr(syn_desc, out_desc)
+    _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)
+        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
@@ -139,6 +143,7 @@ class AlignPostProj(Interaction):
         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__))
 
@@ -154,12 +159,21 @@ class AlignPostProj(Interaction):
         # 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 '
@@ -176,7 +190,11 @@ class AlignPostProj(Interaction):
 
         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)
+            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)
 

brainstate/nn/_elementwise/_dropout_test.py

@@ -59,17 +59,17 @@ class TestDropout(unittest.TestCase):
         expected_non_zero_elements = input_data[output_data != 0] * scale_factor
         np.testing.assert_almost_equal(non_zero_elements, expected_non_zero_elements)
 
-    def test_Dropout1d(self):
-        dropout_layer = bst.nn.Dropout1d(prob=0.5)
-        input_data = np.random.randn(2, 3, 4)
-        with bst.environ.context(fit=True):
-            output_data = dropout_layer(input_data)
-        self.assertEqual(input_data.shape, output_data.shape)
-        self.assertTrue(np.any(output_data == 0))
-        scale_factor = 1 / (1 - 0.5)
-        non_zero_elements = output_data[output_data != 0]
-        expected_non_zero_elements = input_data[output_data != 0] * scale_factor
-        np.testing.assert_almost_equal(non_zero_elements, expected_non_zero_elements, decimal=4)
+    # def test_Dropout1d(self):
+    #     dropout_layer = bst.nn.Dropout1d(prob=0.5)
+    #     input_data = np.random.randn(2, 3, 4)
+    #     with bst.environ.context(fit=True):
+    #         output_data = dropout_layer(input_data)
+    #     self.assertEqual(input_data.shape, output_data.shape)
+    #     self.assertTrue(np.any(output_data == 0))
+    #     scale_factor = 1 / (1 - 0.5)
+    #     non_zero_elements = output_data[output_data != 0]
+    #     expected_non_zero_elements = input_data[output_data != 0] * scale_factor
+    #     np.testing.assert_almost_equal(non_zero_elements, expected_non_zero_elements, decimal=4)
 
     def test_Dropout2d(self):
         dropout_layer = bst.nn.Dropout2d(prob=0.5)

brainstate/nn/_interaction/_linear.py

@@ -20,10 +20,7 @@ from __future__ import annotations
 from typing import Callable, Union, Optional
 
 import brainunit as u
-import jax
 import jax.numpy as jnp
-from jax.experimental.sparse.coo import coo_matvec_p, coo_matmat_p, COOInfo
-from jax.experimental.sparse.csr import csr_matvec_p, csr_matmat_p
 
 from brainstate import init, functional
 from brainstate._state import ParamState
@@ -34,9 +31,7 @@ __all__ = [
     'Linear',
     'ScaledWSLinear',
     'SignedWLinear',
-    'CSRLinear',
-    'CSCLinear',
-    'COOLinear',
+    'SparseLinear',
     'AllToAll',
     'OneToOne',
 ]
@@ -198,270 +193,48 @@ class ScaledWSLinear(Module):
         return y
 
 
-def csr_matmat(data, indices, indptr, B: jax.Array, *, shape, transpose: bool = False) -> jax.Array:
-    """Product of CSR sparse matrix and a dense matrix.
-
-    Args:
-      data : array of shape ``(nse,)``.
-      indices : array of shape ``(nse,)``
-      indptr : array of shape ``(shape[0] + 1,)`` and dtype ``indices.dtype``
-      B : array of shape ``(mat.shape[0] if transpose else mat.shape[1], cols)`` and
-        dtype ``mat.dtype``
-      transpose : boolean specifying whether to transpose the sparse matrix
-        before computing.
-
-    Returns:
-      C : array of shape ``(mat.shape[1] if transpose else mat.shape[0], cols)``
-        representing the matrix vector product.
+class SparseLinear(Module):
     """
-    return csr_matmat_p.bind(data, indices, indptr, B, shape=shape, transpose=transpose)
-
-
-def csr_matvec(data, indices, indptr, v, *, shape, transpose=False) -> jax.Array:
-    """Product of CSR sparse matrix and a dense vector.
+    Linear layer with Sparse Matrix (can be ``brainunit.sparse.CSR``,
+    ``brainunit.sparse.CSC``, ``brainunit.sparse.COO``, or any other sparse matrix).
 
     Args:
-      data : array of shape ``(nse,)``.
-      indices : array of shape ``(nse,)``
-      indptr : array of shape ``(shape[0] + 1,)`` and dtype ``indices.dtype``
-      v : array of shape ``(shape[0] if transpose else shape[1],)``
-        and dtype ``data.dtype``
-      shape : length-2 tuple representing the matrix shape
-      transpose : boolean specifying whether to transpose the sparse matrix
-        before computing.
-
-    Returns:
-      y : array of shape ``(shape[1] if transpose else shape[0],)`` representing
-        the matrix vector product.
-    """
-    return csr_matvec_p.bind(data, indices, indptr, v, shape=shape, transpose=transpose)
-
-
-class CSRLinear(Module):
-    """
-    Linear layer with Compressed Sparse Row (CSR) matrix.
-    """
-    __module__ = 'brainstate.nn'
-
-    def __init__(
-        self,
-        in_size: Size,
-        out_size: Size,
-        indptr: ArrayLike,
-        indices: ArrayLike,
-        weight: Union[Callable, ArrayLike],
-        b_init: Optional[Union[Callable, ArrayLike]] = None,
-        name: Optional[str] = None,
-    ):
-        super().__init__(name=name)
-
-        # input and output shape
-        self.in_size = in_size
-        self.out_size = out_size
-        assert self.in_size[:-1] == self.out_size[:-1], ('The first n-1 dimensions of "in_size" '
-                                                         'and "out_size" must be the same.')
-
-        # 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 == self.in_size[-1] + 1, f"indptr must have size {self.in_size[-1] + 1}. Got: {indptr.size}"
-        with jax.ensure_compile_time_eval():
-            self.indptr = u.math.asarray(indptr)
-            self.indices = u.math.asarray(indices)
-
-        # weights
-        weight = init.param(weight, (len(indices),), allow_none=False, allow_scalar=False)
-        params = dict(weight=weight)
-        if b_init is not None:
-            params['bias'] = init.param(b_init, self.out_size[-1], allow_none=False)
-        self.weight = ParamState(params)
-
-    def update(self, x):
-        data = self.weight.value['weight']
-        data, w_unit = u.get_mantissa(data), u.get_unit(data)
-        x, x_unit = u.get_mantissa(x), u.get_unit(x)
-        shape = [self.in_size[-1], self.out_size[-1]]
-        if x.ndim == 1:
-            y = csr_matvec(data, self.indices, self.indptr, x, shape=shape)
-        elif x.ndim == 2:
-            y = csr_matmat(data, self.indices, self.indptr, x, shape=shape)
-        else:
-            raise NotImplementedError(f"matmul with object of shape {x.shape}")
-        y = u.maybe_decimal(u.Quantity(y, unit=w_unit * x_unit))
-        if 'bias' in self.weight.value:
-            y = y + self.weight.value['bias']
-        return y
-
-
-class CSCLinear(Module):
-    """
-    Linear layer with Compressed Sparse Column (CSC) matrix.
+        spar_mat: SparseMatrix. The sparse weight matrix.
+        in_size: Size. The input size.
+        name: str. The object name.
     """
     __module__ = 'brainstate.nn'
 
     def __init__(
         self,
-        in_size: Size,
-        out_size: Size,
-        indptr: ArrayLike,
-        indices: ArrayLike,
-        weight: Union[Callable, ArrayLike],
-        b_init: Optional[Union[Callable, ArrayLike]] = None,
-        name: Optional[str] = None,
-    ):
-        super().__init__(name=name)
-
-        # input and output shape
-        self.in_size = in_size
-        self.out_size = out_size
-        assert self.in_size[:-1] == self.out_size[:-1], ('The first n-1 dimensions of "in_size" '
-                                                         'and "out_size" must be the same.')
-
-        # 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 == self.in_size[-1] + 1, f"indptr must have size {self.in_size[-1] + 1}. Got: {indptr.size}"
-        with jax.ensure_compile_time_eval():
-            self.indptr = u.math.asarray(indptr)
-            self.indices = u.math.asarray(indices)
-
-        # weights
-        weight = init.param(weight, (len(indices),), allow_none=False, allow_scalar=False)
-        params = dict(weight=weight)
-        if b_init is not None:
-            params['bias'] = init.param(b_init, self.out_size[-1], allow_none=False)
-        self.weight = ParamState(params)
-
-    def update(self, x):
-        data = self.weight.value['weight']
-        data, w_unit = u.get_mantissa(data), u.get_unit(data)
-        x, x_unit = u.get_mantissa(x), u.get_unit(x)
-        shape = [self.out_size[-1], self.in_size[-1]]
-        if x.ndim == 1:
-            y = csr_matvec(data, self.indices, self.indptr, x, shape=shape, transpose=True)
-        elif x.ndim == 2:
-            y = csr_matmat(data, self.indices, self.indptr, x, shape=shape, transpose=True)
-        else:
-            raise NotImplementedError(f"matmul with object of shape {x.shape}")
-        y = u.maybe_decimal(u.Quantity(y, unit=w_unit * x_unit))
-        if 'bias' in self.weight.value:
-            y = y + self.weight.value['bias']
-        return y
-
-
-def coo_matvec(
-    data: jax.Array,
-    row: jax.Array,
-    col: jax.Array,
-    v: jax.Array, *,
-    spinfo: COOInfo,
-    transpose: bool = False
-) -> jax.Array:
-    """Product of COO sparse matrix and a dense vector.
-
-    Args:
-      data : array of shape ``(nse,)``.
-      row : array of shape ``(nse,)``
-      col : array of shape ``(nse,)`` and dtype ``row.dtype``
-      v : array of shape ``(shape[0] if transpose else shape[1],)`` and
-        dtype ``data.dtype``
-      spinfo : COOInfo object containing the shape of the matrix and the dtype
-      transpose : boolean specifying whether to transpose the sparse matrix
-        before computing.
-
-    Returns:
-      y : array of shape ``(shape[1] if transpose else shape[0],)`` representing
-        the matrix vector product.
-    """
-    return coo_matvec_p.bind(data, row, col, v, spinfo=spinfo, transpose=transpose)
-
-
-def coo_matmat(
-    data: jax.Array, row: jax.Array, col: jax.Array, B: jax.Array, *,
-    spinfo: COOInfo, transpose: bool = False
-) -> jax.Array:
-    """Product of COO sparse matrix and a dense matrix.
-
-    Args:
-      data : array of shape ``(nse,)``.
-      row : array of shape ``(nse,)``
-      col : array of shape ``(nse,)`` and dtype ``row.dtype``
-      B : array of shape ``(shape[0] if transpose else shape[1], cols)`` and
-        dtype ``data.dtype``
-      spinfo : COOInfo object containing the shape of the matrix and the dtype
-      transpose : boolean specifying whether to transpose the sparse matrix
-        before computing.
-
-    Returns:
-      C : array of shape ``(shape[1] if transpose else shape[0], cols)``
-        representing the matrix vector product.
-    """
-    return coo_matmat_p.bind(data, row, col, B, spinfo=spinfo, transpose=transpose)
-
-
-class COOLinear(Module):
-
-    def __init__(
-        self,
-        in_size: Size,
-        out_size: Size,
-        row: ArrayLike,
-        col: ArrayLike,
-        weight: Union[Callable, ArrayLike],
+        spar_mat: u.sparse.SparseMatrix,
         b_init: Optional[Union[Callable, ArrayLike]] = None,
-        rows_sorted: bool = False,
-        cols_sorted: bool = False,
+        in_size: Size = None,
         name: Optional[str] = None,
     ):
         super().__init__(name=name)
 
         # input and output shape
-        self.in_size = in_size
-        self.out_size = out_size
-        assert self.in_size[:-1] == self.out_size[:-1], ('The first n-1 dimensions of "in_size" '
-                                                         'and "out_size" must be the same.')
-
-        # COO data structure
-        row = jnp.asarray(row)
-        col = jnp.asarray(col)
-        assert row.ndim == 1, f"row must be 1D. Got: {row.ndim}"
-        assert col.ndim == 1, f"col must be 1D. Got: {col.ndim}"
-        assert row.size == col.size, f"row and col must have the same size. Got: {row.size} and {col.size}"
-        with jax.ensure_compile_time_eval():
-            self.row = u.math.asarray(row)
-            self.col = u.math.asarray(col)
-
-        # COO structure information
-        self.rows_sorted = rows_sorted
-        self.cols_sorted = cols_sorted
+        if in_size is not None:
+            self.in_size = in_size
+        self.out_size = spar_mat.shape[-1]
+        if in_size is not None:
+            assert self.in_size[:-1] == self.out_size[:-1], (
+                'The first n-1 dimensions of "in_size" '
+                'and "out_size" must be the same.'
+            )
 
         # weights
-        weight = init.param(weight, (len(row),), allow_none=False, allow_scalar=False)
-        params = dict(weight=weight)
+        assert isinstance(spar_mat, u.sparse.SparseMatrix), '"weight" must be a SparseMatrix.'
+        self.spar_mat = spar_mat
+        params = dict(weight=spar_mat.data)
         if b_init is not None:
             params['bias'] = init.param(b_init, self.out_size[-1], allow_none=False)
         self.weight = ParamState(params)
 
     def update(self, x):
         data = self.weight.value['weight']
-        data, w_unit = u.get_mantissa(data), u.get_unit(data)
-        x, x_unit = u.get_mantissa(x), u.get_unit(x)
-        spinfo = COOInfo(
-            shape=(self.in_size[-1], self.out_size[-1]),
-            rows_sorted=self.rows_sorted,
-            cols_sorted=self.cols_sorted
-        )
-        if x.ndim == 1:
-            y = coo_matvec(data, self.row, self.col, x, spinfo=spinfo, transpose=False)
-        elif x.ndim == 2:
-            y = coo_matmat(data, self.row, self.col, x, spinfo=spinfo, transpose=False)
-        else:
-            raise NotImplementedError(f"matmul with object of shape {x.shape}")
-        y = u.maybe_decimal(u.Quantity(y, unit=w_unit * x_unit))
+        y = x @ self.spar_mat.with_data(data)
         if 'bias' in self.weight.value:
             y = y + self.weight.value['bias']
         return y