brainstate 0.1.0.post20250105__py2.py3-none-any.whl → 0.1.0.post20250120__py2.py3-none-any.whl

brainstate/nn/_dyn_impl/_inputs.py

@@ -22,8 +22,8 @@ import numpy as np
 
 from brainstate import environ, init, random
 from brainstate._state import ShortTermState
-from brainstate._state import State
-from brainstate.compile import while_loop, cond
+from brainstate._state import State, maybe_state
+from brainstate.compile import while_loop
 from brainstate.nn._dynamics._dynamics_base import Dynamics, Prefetch
 from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike, Size, DTypeLike
@@ -198,55 +198,97 @@ class PoissonInput(Module):
         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)
+        poisson_input(
+            self.freq,
+            self.num_input,
+            self.weight,
+            target_state,
+            self.indices,
         )
 
-        # update target variable
-        target_state.value = target.at[self.indices].add(inp * self.weight)
-
 
 def poisson_input(
-    freq: ArrayLike,
+    freq: u.Quantity[u.Hz],
     num_input: int,
-    weight: ArrayLike,
+    weight: u.Quantity,
     target: State,
     indices: Optional[Union[np.ndarray, jax.Array]] = None,
 ):
     """
     Poisson Input to the given :py:class:`brainstate.State`.
     """
+    freq = maybe_state(freq)
+    weight = maybe_state(weight)
+
     assert isinstance(target, State), 'The target must be a State.'
-    p = freq * environ.get_dt()
+    p = (freq * environ.get_dt()).to_decimal()
     a = num_input * p
     b = num_input * (1 - p)
     tar_val = target.value
+    cond = u.math.logical_and(a > 5, b > 5)
+
     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
+        branch1 = jax.tree.map(
+            lambda tar: random.normal(
+                a,
+                b * p,
+                tar.shape,
+                dtype=tar.dtype
             ),
-            lambda: jax.tree.map(
-                lambda tar: random.binomial(num_input, p, tar.shape).astype(float),
-                tar_val,
-                is_leaf=u.math.is_quantity
-            )
+            tar_val,
+            is_leaf=u.math.is_quantity
+        )
+        branch2 = jax.tree.map(
+            lambda tar: random.binomial(
+                num_input,
+                p,
+                tar.shape,
+                check_valid=False,
+                dtype=tar.dtype
+            ),
+            tar_val,
+            is_leaf=u.math.is_quantity,
+        )
+
+        inp = jax.tree.map(
+            lambda b1, b2: u.math.where(cond, b1, b2),
+            branch1,
+            branch2,
+            is_leaf=u.math.is_quantity,
         )
 
+        # inp = jax.lax.cond(
+        #     cond,
+        #     lambda rand_key: jax.tree.map(
+        #         lambda tar: random.normal(
+        #             a,
+        #             b * p,
+        #             tar.shape,
+        #             key=rand_key,
+        #             dtype=tar.dtype
+        #         ),
+        #         tar_val,
+        #         is_leaf=u.math.is_quantity
+        #     ),
+        #     lambda rand_key: jax.tree.map(
+        #         lambda tar: random.binomial(
+        #             num_input,
+        #             p,
+        #             tar.shape,
+        #             key=rand_key,
+        #             check_valid=False,
+        #             dtype=tar.dtype
+        #         ),
+        #         tar_val,
+        #         is_leaf=u.math.is_quantity,
+        #     ),
+        #     random.split_key()
+        # )
+
         # update target variable
         target.value = jax.tree.map(
             lambda x: x * weight,
@@ -256,19 +298,62 @@ def poisson_input(
 
     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
+        branch1 = jax.tree.map(
+            lambda tar: random.normal(
+                a,
+                b * p,
+                tar[indices].shape,
+                dtype=tar.dtype
             ),
-            lambda: jax.tree.map(
-                lambda tar: random.binomial(num_input, p, tar[indices].shape).astype(float),
-                tar_val,
-                is_leaf=u.math.is_quantity
-            )
+            tar_val,
+            is_leaf=u.math.is_quantity
         )
+        branch2 = jax.tree.map(
+            lambda tar: random.binomial(
+                num_input,
+                p,
+                tar[indices].shape,
+                # check_valid=False,
+                dtype=tar.dtype
+            ),
+            tar_val,
+            is_leaf=u.math.is_quantity
+        )
+
+        inp = jax.tree.map(
+            lambda b1, b2: u.math.where(cond, b1, b2),
+            branch1,
+            branch2,
+            is_leaf=u.math.is_quantity,
+        )
+
+        # inp = jax.lax.cond(
+        #     cond,
+        #     lambda rand_key: jax.tree.map(
+        #         lambda tar: random.normal(
+        #             a,
+        #             b * p,
+        #             tar[indices].shape,
+        #             key=rand_key,
+        #             dtype=tar.dtype
+        #         ),
+        #         tar_val,
+        #         is_leaf=u.math.is_quantity
+        #     ),
+        #     lambda rand_key: jax.tree.map(
+        #         lambda tar: random.binomial(
+        #             num_input,
+        #             p,
+        #             tar[indices].shape,
+        #             key=rand_key,
+        #             check_valid=False,
+        #             dtype=tar.dtype
+        #         ),
+        #         tar_val,
+        #         is_leaf=u.math.is_quantity
+        #     ),
+        #     random.split_key()
+        # )
 
         # update target variable
         target.value = jax.tree.map(

brainstate/nn/_interaction/_conv.py

@@ -191,6 +191,7 @@ class _Conv(_BaseConv):
         b_init: Optional[Union[Callable, ArrayLike]] = None,
         w_mask: Optional[Union[ArrayLike, Callable]] = None,
         name: str = None,
+        param_type: type = ParamState,
     ):
         super().__init__(in_size=in_size,
                          out_channels=out_channels,
@@ -215,7 +216,7 @@ class _Conv(_BaseConv):
             params['bias'] = bias
 
         # The weight operation
-        self.weight = ParamState(params)
+        self.weight = param_type(params)
 
         # Evaluate the output shape
         abstract_y = jax.eval_shape(
@@ -346,6 +347,7 @@ class _ScaledWSConv(_BaseConv):
         b_init: Optional[Union[Callable, ArrayLike]] = None,
         w_mask: Optional[Union[ArrayLike, Callable]] = None,
         name: str = None,
+        param_type: type = ParamState,
     ):
         super().__init__(in_size=in_size,
                          out_channels=out_channels,
@@ -379,7 +381,7 @@ class _ScaledWSConv(_BaseConv):
         self.eps = eps
 
         # The weight operation
-        self.weight = ParamState(params)
+        self.weight = param_type(params)
 
         # Evaluate the output shape
         abstract_y = jax.eval_shape(

brainstate/nn/_interaction/_linear.py

@@ -34,6 +34,7 @@ __all__ = [
     'SparseLinear',
     'AllToAll',
     'OneToOne',
+    'LoRA',
 ]
 
 
@@ -51,6 +52,7 @@ class Linear(Module):
         b_init: Optional[Union[Callable, ArrayLike]] = init.ZeroInit(),
         w_mask: Optional[Union[ArrayLike, Callable]] = None,
         name: Optional[str] = None,
+        param_type: type = ParamState,
     ):
         super().__init__(name=name)
 
@@ -67,7 +69,7 @@ class Linear(Module):
         params = dict(weight=init.param(w_init, (self.in_size[-1], self.out_size[-1]), allow_none=False))
         if b_init is not None:
             params['bias'] = init.param(b_init, self.out_size[-1], allow_none=False)
-        self.weight = ParamState(params)
+        self.weight = param_type(params)
 
     def update(self, x):
         params = self.weight.value
@@ -93,7 +95,7 @@ class SignedWLinear(Module):
         w_init: Union[Callable, ArrayLike] = init.KaimingNormal(),
         w_sign: Optional[ArrayLike] = None,
         name: Optional[str] = None,
-
+        param_type: type = ParamState,
     ):
         super().__init__(name=name)
 
@@ -108,7 +110,7 @@ class SignedWLinear(Module):
 
         # weights
         weight = init.param(w_init, self.in_size + self.out_size, allow_none=False)
-        self.weight = ParamState(weight)
+        self.weight = param_type(weight)
 
     def update(self, x):
         w = self.weight.value
@@ -156,6 +158,7 @@ class ScaledWSLinear(Module):
         ws_gain: bool = True,
         eps: float = 1e-4,
         name: str = None,
+        param_type: type = ParamState,
     ):
         super().__init__(name=name)
 
@@ -179,7 +182,7 @@ class ScaledWSLinear(Module):
         if ws_gain:
             s = params['weight'].shape
             params['gain'] = jnp.ones((1,) * (len(s) - 1) + (s[-1],), dtype=params['weight'].dtype)
-        self.weight = ParamState(params)
+        self.weight = param_type(params)
 
     def update(self, x):
         params = self.weight.value
@@ -211,6 +214,7 @@ class SparseLinear(Module):
         b_init: Optional[Union[Callable, ArrayLike]] = None,
         in_size: Size = None,
         name: Optional[str] = None,
+        param_type: type = ParamState,
     ):
         super().__init__(name=name)
 
@@ -230,7 +234,7 @@ class SparseLinear(Module):
         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)
+        self.weight = param_type(params)
 
     def update(self, x):
         data = self.weight.value['weight']
@@ -260,6 +264,7 @@ class AllToAll(Module):
         b_init: Optional[Union[Callable, ArrayLike]] = None,
         include_self: bool = True,
         name: Optional[str] = None,
+        param_type: type = ParamState,
     ):
         super().__init__(name=name)
 
@@ -277,7 +282,7 @@ class AllToAll(Module):
         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)
+        self.weight = param_type(params)
 
     def update(self, pre_val):
         params = self.weight.value
@@ -332,6 +337,7 @@ class OneToOne(Module):
         w_init: Union[Callable, ArrayLike] = init.Normal(),
         b_init: Optional[Union[Callable, ArrayLike]] = None,
         name: Optional[str] = None,
+        param_type: type = ParamState,
     ):
         super().__init__(name=name)
 
@@ -343,13 +349,81 @@ class OneToOne(Module):
         param = dict(weight=init.param(w_init, self.in_size, allow_none=False))
         if b_init is not None:
             param['bias'] = init.param(b_init, self.out_size, allow_none=False)
-        self.weight = param
+        self.weight = param_type(param)
 
     def update(self, pre_val):
         pre_val, pre_unit = u.get_mantissa(pre_val), u.get_unit(pre_val)
-        w_val, w_unit = u.get_mantissa(self.weight['weight']), u.get_unit(self.weight['weight'])
+        w_val, w_unit = u.get_mantissa(self.weight.value['weight']), u.get_unit(self.weight.value['weight'])
         post_val = pre_val * w_val
         post_val = u.maybe_decimal(u.Quantity(post_val, unit=w_unit * pre_unit))
-        if 'bias' in self.weight:
-            post_val = post_val + self.weight['bias']
+        if 'bias' in self.weight.value:
+            post_val = post_val + self.weight.value['bias']
         return post_val
+
+
+class LoRA(Module):
+    """A standalone LoRA layer.
+
+    Example usage::
+
+        >>> import brainstate as bst
+        >>> import jax, jax.numpy as jnp
+        >>> layer = bst.nn.LoRA(3, 2, 4)
+        >>> layer.weight.value
+    {'lora_a': Array([[ 0.25141352, -0.09826107],
+            [ 0.2328382 ,  0.38869813],
+            [ 0.27069277,  0.7678282 ]], dtype=float32),
+     'lora_b': Array([[-0.8372317 ,  0.21012013, -0.52999765, -0.31939325],
+            [ 0.64234126, -0.42980042,  1.2549229 , -0.47134295]],      dtype=float32)}
+        >>> # Wrap around existing layer
+        >>> linear = bst.nn.Linear(3, 4)
+        >>> wrapper = bst.nn.LoRA(3, 2, 4, base_module=linear)
+        >>> assert wrapper.base_module == linear
+        >>> y = layer(jnp.ones((16, 3)))
+        >>> y.shape
+        (16, 4)
+
+    Args:
+        in_features: the number of input features.
+        lora_rank: the rank of the LoRA dimension.
+        out_features: the number of output features.
+        base_module: a base module to call and substitute, if possible.
+        kernel_init: initializer function for the weight matrices.
+        param_type: the type of the LoRA params.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        lora_rank: int,
+        out_features: int,
+        *,
+        base_module: Optional[Module] = None,
+        kernel_init: Union[Callable, ArrayLike] = init.LecunNormal(),
+        param_type: type = ParamState,
+    ):
+        super().__init__()
+
+        # input and output shape
+        self.in_size = in_features
+        self.out_size = out_features
+        self.in_features = in_features
+        self.out_features = out_features
+
+        # others
+        self.base_module = base_module
+
+        # weights
+        param = dict(
+            lora_a=kernel_init((in_features, lora_rank)),
+            lora_b=kernel_init((lora_rank, out_features))
+        )
+        self.weight = param_type(param)
+
+    def __call__(self, x: ArrayLike):
+        out = x @ self.weight.value['lora_a'] @ self.weight.value['lora_b']
+        if self.base_module is not None:
+            if not callable(self.base_module):
+                raise ValueError('`self.base_module` must be callable.')
+            out += self.base_module(x)
+        return out

brainstate/random/_rand_funs.py

@@ -1848,7 +1848,14 @@ def lognormal(mean=None, sigma=None, size: Optional[Size] = None,
     return DEFAULT.lognormal(mean, sigma, size, key=key, dtype=dtype)
 
 
-def binomial(n, p, size: Optional[Size] = None, key: Optional[SeedOrKey] = None, dtype: DTypeLike = None):
+def binomial(
+    n,
+    p,
+    size: Optional[Size] = None,
+    key: Optional[SeedOrKey] = None,
+    dtype: DTypeLike = None,
+    check_valid: bool = True,
+):
     r"""
     Draw samples from a binomial distribution.
 
@@ -1933,7 +1940,7 @@ def binomial(n, p, size: Optional[Size] = None, key: Optional[SeedOrKey] = None,
     >>> sum(brainstate.random.binomial(9, 0.1, 20000) == 0)/20000.
     # answer = 0.38885, or 38%.
     """
-    return DEFAULT.binomial(n, p, size, key=key, dtype=dtype)
+    return DEFAULT.binomial(n, p, size, key=key, dtype=dtype, check_valid=check_valid)
 
 
 def chisquare(df, size: Optional[Size] = None, key: Optional[SeedOrKey] = None, dtype: DTypeLike = None):

brainstate/random/_rand_seed.py

@@ -21,7 +21,7 @@ import jax
 import numpy as np
 
 from brainstate.typing import SeedOrKey
-from ._rand_state import RandomState, DEFAULT
+from ._rand_state import RandomState, DEFAULT, use_prng_key
 
 __all__ = [
     'seed', 'set_key', 'get_key', 'default_rng', 'split_key', 'split_keys', 'seed_context', 'restore_key',
@@ -123,7 +123,17 @@ def set_key(seed_or_key: SeedOrKey):
     seed_or_key: int
       The random key.
     """
-    DEFAULT.set_key(jax.random.PRNGKey(seed_or_key) if jax.numpy.shape(seed_or_key) == () else seed_or_key)
+    if isinstance(seed_or_key, int):
+        # key = jax.random.key(seed_or_key)
+        key = jax.random.PRNGKey(seed_or_key) if use_prng_key else jrjax.random.key(seed_or_key)
+    elif isinstance(seed_or_key, (jax.numpy.ndarray, np.ndarray)):
+        if jax.numpy.issubdtype(seed_or_key.dtype, jax.dtypes.prng_key):
+            key = seed_or_key
+        elif seed_or_key.size == 2 and seed_or_key.dtype == jax.numpy.uint32:
+            key = seed_or_key
+        else:
+            raise ValueError(f"seed_or_key should be an integer or a tuple of two integers.")
+    DEFAULT.set_key(key)
 
 
 def get_key():