brainstate 0.1.8__py2.py3-none-any.whl → 0.1.9__py2.py3-none-any.whl

brainstate/nn/_elementwise_test.py

@@ -1,169 +1,169 @@
-# 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 absl.testing import absltest
-from absl.testing import parameterized
-
-import brainstate
-
-
-class Test_Activation(parameterized.TestCase):
-
-    def test_Threshold(self):
-        threshold_layer = brainstate.nn.Threshold(5, 20)
-        input = brainstate.random.randn(2)
-        output = threshold_layer(input)
-
-    def test_ReLU(self):
-        ReLU_layer = brainstate.nn.ReLU()
-        input = brainstate.random.randn(2)
-        output = ReLU_layer(input)
-
-    def test_RReLU(self):
-        RReLU_layer = brainstate.nn.RReLU(lower=0, upper=1)
-        input = brainstate.random.randn(2)
-        output = RReLU_layer(input)
-
-    def test_Hardtanh(self):
-        Hardtanh_layer = brainstate.nn.Hardtanh(min_val=0, max_val=1, )
-        input = brainstate.random.randn(2)
-        output = Hardtanh_layer(input)
-
-    def test_ReLU6(self):
-        ReLU6_layer = brainstate.nn.ReLU6()
-        input = brainstate.random.randn(2)
-        output = ReLU6_layer(input)
-
-    def test_Sigmoid(self):
-        Sigmoid_layer = brainstate.nn.Sigmoid()
-        input = brainstate.random.randn(2)
-        output = Sigmoid_layer(input)
-
-    def test_Hardsigmoid(self):
-        Hardsigmoid_layer = brainstate.nn.Hardsigmoid()
-        input = brainstate.random.randn(2)
-        output = Hardsigmoid_layer(input)
-
-    def test_Tanh(self):
-        Tanh_layer = brainstate.nn.Tanh()
-        input = brainstate.random.randn(2)
-        output = Tanh_layer(input)
-
-    def test_SiLU(self):
-        SiLU_layer = brainstate.nn.SiLU()
-        input = brainstate.random.randn(2)
-        output = SiLU_layer(input)
-
-    def test_Mish(self):
-        Mish_layer = brainstate.nn.Mish()
-        input = brainstate.random.randn(2)
-        output = Mish_layer(input)
-
-    def test_Hardswish(self):
-        Hardswish_layer = brainstate.nn.Hardswish()
-        input = brainstate.random.randn(2)
-        output = Hardswish_layer(input)
-
-    def test_ELU(self):
-        ELU_layer = brainstate.nn.ELU(alpha=0.5, )
-        input = brainstate.random.randn(2)
-        output = ELU_layer(input)
-
-    def test_CELU(self):
-        CELU_layer = brainstate.nn.CELU(alpha=0.5, )
-        input = brainstate.random.randn(2)
-        output = CELU_layer(input)
-
-    def test_SELU(self):
-        SELU_layer = brainstate.nn.SELU()
-        input = brainstate.random.randn(2)
-        output = SELU_layer(input)
-
-    def test_GLU(self):
-        GLU_layer = brainstate.nn.GLU()
-        input = brainstate.random.randn(4, 2)
-        output = GLU_layer(input)
-
-    @parameterized.product(
-        approximate=['tanh', 'none']
-    )
-    def test_GELU(self, approximate):
-        GELU_layer = brainstate.nn.GELU()
-        input = brainstate.random.randn(2)
-        output = GELU_layer(input)
-
-    def test_Hardshrink(self):
-        Hardshrink_layer = brainstate.nn.Hardshrink(lambd=1)
-        input = brainstate.random.randn(2)
-        output = Hardshrink_layer(input)
-
-    def test_LeakyReLU(self):
-        LeakyReLU_layer = brainstate.nn.LeakyReLU()
-        input = brainstate.random.randn(2)
-        output = LeakyReLU_layer(input)
-
-    def test_LogSigmoid(self):
-        LogSigmoid_layer = brainstate.nn.LogSigmoid()
-        input = brainstate.random.randn(2)
-        output = LogSigmoid_layer(input)
-
-    def test_Softplus(self):
-        Softplus_layer = brainstate.nn.Softplus()
-        input = brainstate.random.randn(2)
-        output = Softplus_layer(input)
-
-    def test_Softshrink(self):
-        Softshrink_layer = brainstate.nn.Softshrink(lambd=1)
-        input = brainstate.random.randn(2)
-        output = Softshrink_layer(input)
-
-    def test_PReLU(self):
-        PReLU_layer = brainstate.nn.PReLU(num_parameters=2, init=0.5)
-        input = brainstate.random.randn(2)
-        output = PReLU_layer(input)
-
-    def test_Softsign(self):
-        Softsign_layer = brainstate.nn.Softsign()
-        input = brainstate.random.randn(2)
-        output = Softsign_layer(input)
-
-    def test_Tanhshrink(self):
-        Tanhshrink_layer = brainstate.nn.Tanhshrink()
-        input = brainstate.random.randn(2)
-        output = Tanhshrink_layer(input)
-
-    def test_Softmin(self):
-        Softmin_layer = brainstate.nn.Softmin(dim=2)
-        input = brainstate.random.randn(2, 3, 4)
-        output = Softmin_layer(input)
-
-    def test_Softmax(self):
-        Softmax_layer = brainstate.nn.Softmax(dim=2)
-        input = brainstate.random.randn(2, 3, 4)
-        output = Softmax_layer(input)
-
-    def test_Softmax2d(self):
-        Softmax2d_layer = brainstate.nn.Softmax2d()
-        input = brainstate.random.randn(2, 3, 12, 13)
-        output = Softmax2d_layer(input)
-
-    def test_LogSoftmax(self):
-        LogSoftmax_layer = brainstate.nn.LogSoftmax(dim=2)
-        input = brainstate.random.randn(2, 3, 4)
-        output = LogSoftmax_layer(input)
-
-
-if __name__ == '__main__':
-    absltest.main()
+# 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 absl.testing import absltest
+from absl.testing import parameterized
+
+import brainstate
+
+
+class Test_Activation(parameterized.TestCase):
+
+    def test_Threshold(self):
+        threshold_layer = brainstate.nn.Threshold(5, 20)
+        input = brainstate.random.randn(2)
+        output = threshold_layer(input)
+
+    def test_ReLU(self):
+        ReLU_layer = brainstate.nn.ReLU()
+        input = brainstate.random.randn(2)
+        output = ReLU_layer(input)
+
+    def test_RReLU(self):
+        RReLU_layer = brainstate.nn.RReLU(lower=0, upper=1)
+        input = brainstate.random.randn(2)
+        output = RReLU_layer(input)
+
+    def test_Hardtanh(self):
+        Hardtanh_layer = brainstate.nn.Hardtanh(min_val=0, max_val=1, )
+        input = brainstate.random.randn(2)
+        output = Hardtanh_layer(input)
+
+    def test_ReLU6(self):
+        ReLU6_layer = brainstate.nn.ReLU6()
+        input = brainstate.random.randn(2)
+        output = ReLU6_layer(input)
+
+    def test_Sigmoid(self):
+        Sigmoid_layer = brainstate.nn.Sigmoid()
+        input = brainstate.random.randn(2)
+        output = Sigmoid_layer(input)
+
+    def test_Hardsigmoid(self):
+        Hardsigmoid_layer = brainstate.nn.Hardsigmoid()
+        input = brainstate.random.randn(2)
+        output = Hardsigmoid_layer(input)
+
+    def test_Tanh(self):
+        Tanh_layer = brainstate.nn.Tanh()
+        input = brainstate.random.randn(2)
+        output = Tanh_layer(input)
+
+    def test_SiLU(self):
+        SiLU_layer = brainstate.nn.SiLU()
+        input = brainstate.random.randn(2)
+        output = SiLU_layer(input)
+
+    def test_Mish(self):
+        Mish_layer = brainstate.nn.Mish()
+        input = brainstate.random.randn(2)
+        output = Mish_layer(input)
+
+    def test_Hardswish(self):
+        Hardswish_layer = brainstate.nn.Hardswish()
+        input = brainstate.random.randn(2)
+        output = Hardswish_layer(input)
+
+    def test_ELU(self):
+        ELU_layer = brainstate.nn.ELU(alpha=0.5, )
+        input = brainstate.random.randn(2)
+        output = ELU_layer(input)
+
+    def test_CELU(self):
+        CELU_layer = brainstate.nn.CELU(alpha=0.5, )
+        input = brainstate.random.randn(2)
+        output = CELU_layer(input)
+
+    def test_SELU(self):
+        SELU_layer = brainstate.nn.SELU()
+        input = brainstate.random.randn(2)
+        output = SELU_layer(input)
+
+    def test_GLU(self):
+        GLU_layer = brainstate.nn.GLU()
+        input = brainstate.random.randn(4, 2)
+        output = GLU_layer(input)
+
+    @parameterized.product(
+        approximate=['tanh', 'none']
+    )
+    def test_GELU(self, approximate):
+        GELU_layer = brainstate.nn.GELU()
+        input = brainstate.random.randn(2)
+        output = GELU_layer(input)
+
+    def test_Hardshrink(self):
+        Hardshrink_layer = brainstate.nn.Hardshrink(lambd=1)
+        input = brainstate.random.randn(2)
+        output = Hardshrink_layer(input)
+
+    def test_LeakyReLU(self):
+        LeakyReLU_layer = brainstate.nn.LeakyReLU()
+        input = brainstate.random.randn(2)
+        output = LeakyReLU_layer(input)
+
+    def test_LogSigmoid(self):
+        LogSigmoid_layer = brainstate.nn.LogSigmoid()
+        input = brainstate.random.randn(2)
+        output = LogSigmoid_layer(input)
+
+    def test_Softplus(self):
+        Softplus_layer = brainstate.nn.Softplus()
+        input = brainstate.random.randn(2)
+        output = Softplus_layer(input)
+
+    def test_Softshrink(self):
+        Softshrink_layer = brainstate.nn.Softshrink(lambd=1)
+        input = brainstate.random.randn(2)
+        output = Softshrink_layer(input)
+
+    def test_PReLU(self):
+        PReLU_layer = brainstate.nn.PReLU(num_parameters=2, init=0.5)
+        input = brainstate.random.randn(2)
+        output = PReLU_layer(input)
+
+    def test_Softsign(self):
+        Softsign_layer = brainstate.nn.Softsign()
+        input = brainstate.random.randn(2)
+        output = Softsign_layer(input)
+
+    def test_Tanhshrink(self):
+        Tanhshrink_layer = brainstate.nn.Tanhshrink()
+        input = brainstate.random.randn(2)
+        output = Tanhshrink_layer(input)
+
+    def test_Softmin(self):
+        Softmin_layer = brainstate.nn.Softmin(dim=2)
+        input = brainstate.random.randn(2, 3, 4)
+        output = Softmin_layer(input)
+
+    def test_Softmax(self):
+        Softmax_layer = brainstate.nn.Softmax(dim=2)
+        input = brainstate.random.randn(2, 3, 4)
+        output = Softmax_layer(input)
+
+    def test_Softmax2d(self):
+        Softmax2d_layer = brainstate.nn.Softmax2d()
+        input = brainstate.random.randn(2, 3, 12, 13)
+        output = Softmax2d_layer(input)
+
+    def test_LogSoftmax(self):
+        LogSoftmax_layer = brainstate.nn.LogSoftmax(dim=2)
+        input = brainstate.random.randn(2, 3, 4)
+        output = LogSoftmax_layer(input)
+
+
+if __name__ == '__main__':
+    absltest.main()

brainstate/nn/_embedding.py

@@ -1,58 +1,58 @@
-# 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 Optional, Callable, Union
-
-from brainstate import init
-from brainstate._state import ParamState
-from brainstate.typing import ArrayLike
-from ._module import Module
-
-__all__ = [
-    'Embedding',
-]
-
-
-class Embedding(Module):
-    r"""
-    A simple lookup table that stores embeddings of a fixed size.
-
-    Args:
-      num_embeddings: Size of embedding dictionary. Must be non-negative.
-      embedding_size: Size of each embedding vector. Must be non-negative.
-      embedding_init: The initializer for the embedding lookup table, of shape `(num_embeddings, embedding_size)`.
-    """
-
-    def __init__(
-        self,
-        num_embeddings: int,
-        embedding_size: int,
-        embedding_init: Union[Callable, ArrayLike] = init.LecunUniform(),
-        name: Optional[str] = None,
-    ):
-        super().__init__(name=name)
-        if num_embeddings < 0:
-            raise ValueError("num_embeddings must not be negative.")
-        if embedding_size < 0:
-            raise ValueError("embedding_size must not be negative.")
-        self.num_embeddings = num_embeddings
-        self.embedding_size = embedding_size
-        self.out_size = (embedding_size,)
-
-        weight = init.param(embedding_init, (self.num_embeddings, self.embedding_size))
-        self.weight = ParamState(weight)
-
-    def update(self, indices: ArrayLike):
-        return self.weight.value[indices]
+# 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 Optional, Callable, Union
+
+from brainstate import init
+from brainstate._state import ParamState
+from brainstate.typing import ArrayLike
+from ._module import Module
+
+__all__ = [
+    'Embedding',
+]
+
+
+class Embedding(Module):
+    r"""
+    A simple lookup table that stores embeddings of a fixed size.
+
+    Args:
+      num_embeddings: Size of embedding dictionary. Must be non-negative.
+      embedding_size: Size of each embedding vector. Must be non-negative.
+      embedding_init: The initializer for the embedding lookup table, of shape `(num_embeddings, embedding_size)`.
+    """
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_size: int,
+        embedding_init: Union[Callable, ArrayLike] = init.LecunUniform(),
+        name: Optional[str] = None,
+    ):
+        super().__init__(name=name)
+        if num_embeddings < 0:
+            raise ValueError("num_embeddings must not be negative.")
+        if embedding_size < 0:
+            raise ValueError("embedding_size must not be negative.")
+        self.num_embeddings = num_embeddings
+        self.embedding_size = embedding_size
+        self.out_size = (embedding_size,)
+
+        weight = init.param(embedding_init, (self.num_embeddings, self.embedding_size))
+        self.weight = ParamState(weight)
+
+    def update(self, indices: ArrayLike):
+        return self.weight.value[indices]

brainstate/nn/_exp_euler.py

@@ -1,92 +1,92 @@
-# 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 Callable
-
-import brainunit as u
-import jax.numpy as jnp
-
-from brainstate import environ, random
-from brainstate.augment import vector_grad
-
-__all__ = [
-    'exp_euler_step',
-]
-
-
-def exp_euler_step(
-    fn: Callable, *args, **kwargs
-):
-    r"""
-    One-step Exponential Euler method for solving ODEs.
-
-    Examples
-    --------
-
-    >>> def fun(x, t):
-    ...     return -x
-    >>> x = 1.0
-    >>> exp_euler_step(fun, x, None)
-
-    If the variable ( $x$ ) has units of ( $[X]$ ), then the drift term ( $\text{drift_fn}(x)$ ) should
-    have units of ( $[X]/[T]$ ), where ( $[T]$ ) is the unit of time.
-
-    If the variable ( x ) has units of ( [X] ), then the diffusion term ( \text{diffusion_fn}(x) )
-    should have units of ( [X]/\sqrt{[T]} ).
-
-    Args:
-        fun: Callable. The function to be solved.
-        diffusion: Callable. The diffusion function.
-        *args: The input arguments.
-        drift: Callable. The drift function.
-
-    Returns:
-        The one-step solution of the ODE.
-    """
-    assert callable(fn), 'The input function should be callable.'
-    assert len(args) > 0, 'The input arguments should not be empty.'
-    if callable(args[0]):
-        diffusion = args[0]
-        args = args[1:]
-    else:
-        diffusion = None
-    assert len(args) > 0, 'The input arguments should not be empty.'
-    if u.math.get_dtype(args[0]) not in [jnp.float32, jnp.float64, jnp.float16, jnp.bfloat16]:
-        raise ValueError(
-            f'The input data type should be float64, float32, float16, or bfloat16 '
-            f'when using Exponential Euler method. But we got {args[0].dtype}.'
-        )
-
-    # drift
-    dt = environ.get('dt')
-    linear, derivative = vector_grad(fn, argnums=0, return_value=True)(*args, **kwargs)
-    linear = u.Quantity(u.get_mantissa(linear), u.get_unit(derivative) / u.get_unit(linear))
-    phi = u.math.exprel(dt * linear)
-    x_next = args[0] + dt * phi * derivative
-
-    # diffusion
-    if diffusion is not None:
-        diffusion_part = diffusion(*args, **kwargs) * u.math.sqrt(dt) * random.randn_like(args[0])
-        if u.get_dim(x_next) != u.get_dim(diffusion_part):
-            drift_unit = u.get_unit(x_next)
-            time_unit = u.get_unit(dt)
-            raise ValueError(
-                f"Drift unit is {drift_unit}, "
-                f"expected diffusion unit is {drift_unit / time_unit ** 0.5}, "
-                f"but we got {u.get_unit(diffusion_part)}."
-            )
-        x_next += diffusion_part
-    return x_next
+# 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 Callable
+
+import brainunit as u
+import jax.numpy as jnp
+
+from brainstate import environ, random
+from brainstate.augment import vector_grad
+
+__all__ = [
+    'exp_euler_step',
+]
+
+
+def exp_euler_step(
+    fn: Callable, *args, **kwargs
+):
+    r"""
+    One-step Exponential Euler method for solving ODEs.
+
+    Examples
+    --------
+
+    >>> def fun(x, t):
+    ...     return -x
+    >>> x = 1.0
+    >>> exp_euler_step(fun, x, None)
+
+    If the variable ( $x$ ) has units of ( $[X]$ ), then the drift term ( $\text{drift_fn}(x)$ ) should
+    have units of ( $[X]/[T]$ ), where ( $[T]$ ) is the unit of time.
+
+    If the variable ( x ) has units of ( [X] ), then the diffusion term ( \text{diffusion_fn}(x) )
+    should have units of ( [X]/\sqrt{[T]} ).
+
+    Args:
+        fun: Callable. The function to be solved.
+        diffusion: Callable. The diffusion function.
+        *args: The input arguments.
+        drift: Callable. The drift function.
+
+    Returns:
+        The one-step solution of the ODE.
+    """
+    assert callable(fn), 'The input function should be callable.'
+    assert len(args) > 0, 'The input arguments should not be empty.'
+    if callable(args[0]):
+        diffusion = args[0]
+        args = args[1:]
+    else:
+        diffusion = None
+    assert len(args) > 0, 'The input arguments should not be empty.'
+    if u.math.get_dtype(args[0]) not in [jnp.float32, jnp.float64, jnp.float16, jnp.bfloat16]:
+        raise ValueError(
+            f'The input data type should be float64, float32, float16, or bfloat16 '
+            f'when using Exponential Euler method. But we got {args[0].dtype}.'
+        )
+
+    # drift
+    dt = environ.get('dt')
+    linear, derivative = vector_grad(fn, argnums=0, return_value=True)(*args, **kwargs)
+    linear = u.Quantity(u.get_mantissa(linear), u.get_unit(derivative) / u.get_unit(linear))
+    phi = u.math.exprel(dt * linear)
+    x_next = args[0] + dt * phi * derivative
+
+    # diffusion
+    if diffusion is not None:
+        diffusion_part = diffusion(*args, **kwargs) * u.math.sqrt(dt) * random.randn_like(args[0])
+        if u.get_dim(x_next) != u.get_dim(diffusion_part):
+            drift_unit = u.get_unit(x_next)
+            time_unit = u.get_unit(dt)
+            raise ValueError(
+                f"Drift unit is {drift_unit}, "
+                f"expected diffusion unit is {drift_unit / time_unit ** 0.5}, "
+                f"but we got {u.get_unit(diffusion_part)}."
+            )
+        x_next += diffusion_part
+    return x_next