Update documentation and examples

.gitignore

@@ -229,3 +229,7 @@ cython_debug/
 /bugs/
 /dev/
 /.claude/
+/docs_version2/_static/logos/
+/docs/_static/logos/
+/docs/_build/
+/docs/changelog.md

.readthedocs.yml

@@ -13,7 +13,7 @@ build:
     post_checkout:
       # Copy the appropriate conf.py based on project name
       - |
-        if [ "$READTHEDOCS_PROJECT" = "brainpy-version2" ]; then
+        if [ "$PROJECT_VERSION" = "brainpy-version2" ]; then
           mkdir -p docs_build
           cp -r docs_version2/* docs_build/
         else

README.md

@@ -74,7 +74,7 @@ We provide a Binder environment for BrainPy. You can use the following button to
 
 ## Citing 
 
-If you are using ``brainpy``, please consider citing the corresponding paper:
+If you are using ``brainpy >= 2.0``, please consider citing the corresponding paper:
 
 ```bibtex
 @article {10.7554/eLife.86365,

brainpy/__init__.py

@@ -16,6 +16,7 @@
 __version__ = "3.0.0"
 __version_info__ = (3, 0, 0)
 
+from . import mixin
 from . import version2
 from ._base import *
 from ._base import __all__ as base_all
@@ -40,7 +41,7 @@
 from ._synouts import *
 from ._synouts import __all__ as synout_all
 
-__main__ = ['version2'] + errors_all + inputs_all + neuron_all + readout_all + stp_all + synapse_all
+__main__ = ['version2', 'mixin'] + errors_all + inputs_all + neuron_all + readout_all + stp_all + synapse_all
 __main__ = __main__ + synout_all + base_all + exp_all + proj_all + synproj_all
 del errors_all, inputs_all, neuron_all, readout_all, stp_all, synapse_all, synout_all, base_all
 del exp_all, proj_all, synproj_all
@@ -81,7 +82,7 @@ def __dir__():
 _deprecated_modules = [
     'math', 'check', 'tools', 'connect', 'initialize', 'init', 'conn',
     'optim', 'losses', 'measure', 'inputs', 'encoding', 'checkpoints',
-    'mixin', 'algorithms', 'integrators', 'ode', 'sde', 'fde',
+    'algorithms', 'integrators', 'ode', 'sde', 'fde',
     'dnn', 'layers', 'dyn', 'running', 'train', 'analysis',
     'channels', 'neurons', 'rates', 'synapses', 'synouts', 'synplast',
     'visualization', 'visualize', 'types', 'modes', 'context',
@@ -96,4 +97,3 @@ def __dir__():
 del _sys, _mod_name, _deprecated_modules
 
 version2.__version__ = __version__
-

brainpy/_base.py

@@ -16,6 +16,7 @@
 from typing import Callable, Optional
 
 import braintools
+
 import brainstate
 
 __all__ = [
@@ -43,12 +44,12 @@ class Neuron(brainstate.nn.Dynamics):
         for multi-dimensional input (e.g., ``100`` or ``(28, 28)``).
     spk_fun : Callable, optional
         Surrogate gradient function for the non-differentiable spike generation operation.
-        Default is ``brainstate.surrogate.InvSquareGrad()``. Common alternatives include:
+        Default is ``braintools.surrogate.InvSquareGrad()``. Common alternatives include:
 
-        - ``brainstate.surrogate.ReluGrad()``
-        - ``brainstate.surrogate.SigmoidGrad()``
-        - ``brainstate.surrogate.GaussianGrad()``
-        - ``brainstate.surrogate.ATan()``
+        - ``braintools.surrogate.ReluGrad()``
+        - ``braintools.surrogate.SigmoidGrad()``
+        - ``braintools.surrogate.GaussianGrad()``
+        - ``braintools.surrogate.ATan()``
     spk_reset : str, optional
         Reset mechanism applied after spike generation. Default is ``'soft'``.
 
@@ -149,7 +150,7 @@ class Neuron(brainstate.nn.Dynamics):
         ...     in_size=100,
         ...     tau=10*u.ms,
         ...     V_th=1.0*u.mV,
-        ...     spk_fun=brainstate.surrogate.ReluGrad(),
+        ...     spk_fun=braintools.surrogate.ReluGrad(),
         ...     spk_reset='soft'
         ... )
         >>>
@@ -304,7 +305,7 @@ class Synapse(brainstate.nn.Dynamics):
 
     **Alignment Patterns**
 
-    Some synapse models inherit from ``brainstate.mixin.AlignPost`` to enable
+    Some synapse models inherit from :class:`AlignPost` to enable
     event-driven computation where synaptic variables are aligned with postsynaptic
     neurons. This is particularly efficient for sparse connectivity patterns.
 

brainpy/_base_test.py

@@ -1,4 +1,4 @@
-# Copyright 2025 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2025 BrainX 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.
@@ -29,12 +29,11 @@
 
 import unittest
 
-import brainstate
 import braintools
 import brainunit as u
-import jax
 import jax.numpy as jnp
 
+import brainstate
 from brainpy._base import Neuron, Synapse
 
 

brainpy/_errors.py

@@ -33,7 +33,6 @@
     'MonitorError',
     'MathError',
     'JaxTracerError',
-    'ConcretizationTypeError',
     'GPUOperatorNotFound',
     'SharedArgError',
 ]
@@ -128,120 +127,6 @@ class MathError(BrainPyError):
     __module__ = 'brainpy'
 
 
-class MPACheckpointingRequiredError(BrainPyError):
-    """To optimally save and restore a multiprocess array (GDA or jax Array outputted from pjit), use GlobalAsyncCheckpointManager.
-
-    You can create an GlobalAsyncCheckpointManager at top-level and pass it as
-    argument::
-
-      from jax.experimental.gda_serialization import serialization as gdas
-      gda_manager = gdas.GlobalAsyncCheckpointManager()
-      brainpy.checkpoints.save(..., gda_manager=gda_manager)
-    """
-    __module__ = 'brainpy'
-
-    def __init__(self, path, step):
-        super().__init__(
-            f'Checkpoint failed at step: "{step}" and path: "{path}": Target '
-            'contains a multiprocess array should be saved/restored with a '
-            'GlobalAsyncCheckpointManager.')
-
-
-class MPARestoreTargetRequiredError(BrainPyError):
-    """Provide a valid target when restoring a checkpoint with a multiprocess array.
-
-    Multiprocess arrays need a sharding (global meshes and partition specs) to be
-    initialized. Therefore, to restore a checkpoint that contains a multiprocess
-    array, make sure the ``target`` you passed contains valid multiprocess arrays
-    at the corresponding tree structure location. If you cannot provide a full
-    valid ``target``, consider ``allow_partial_mpa_restoration=True``.
-    """
-    __module__ = 'brainpy'
-
-    def __init__(self, path, step, key=None):
-        error_msg = (
-            f'Restore checkpoint failed at step: "{step}" and path: "{path}": '
-            'Checkpoints containing a multiprocess array need to be restored with '
-            'a target with pre-created arrays. If you cannot provide a full valid '
-            'target, consider ``allow_partial_mpa_restoration=True``. ')
-        if key:
-            error_msg += f'This error fired when trying to restore array at {key}.'
-        super().__init__(error_msg)
-
-
-class MPARestoreDataCorruptedError(BrainPyError):
-    """A multiprocess array stored in Google Cloud Storage doesn't contain a "commit_success.txt" file, which should be written at the end of the save.
-
-    Failure of finding it could indicate a corruption of your saved GDA data.
-    """
-    __module__ = 'brainpy'
-
-    def __init__(self, step, path):
-        super().__init__(
-            f'Restore checkpoint failed at step: "{step}" on multiprocess array at '
-            f' "{path}": No "commit_success.txt" found on this "_gda" directory. '
-            'Was its save halted before completion?')
-
-
-class MPARestoreTypeNotMatchError(BrainPyError):
-    """Make sure the multiprocess array type you use matches your configuration in jax.config.jax_array.
-
-    If you turned `jax.config.jax_array` on, you should use
-    `jax.experimental.array.Array` everywhere, instead of using
-    `GlobalDeviceArray`. Otherwise, avoid using jax.experimental.array
-    to restore your checkpoint.
-    """
-    __module__ = 'brainpy'
-
-    def __init__(self, step, gda_path):
-        super().__init__(
-            f'Restore checkpoint failed at step: "{step}" on multiprocess array at '
-            f' "{gda_path}": The array type provided by the target does not match '
-            'the JAX global configuration, namely the jax.config.jax_array.')
-
-
-class AlreadyExistsError(BrainPyError):
-    """Attempting to overwrite a file via copy.
-
-    You can pass ``overwrite=True`` to disable this behavior and overwrite
-    existing files in.
-    """
-    __module__ = 'brainpy'
-
-    def __init__(self, path):
-        super().__init__(f'Trying overwrite an existing file: "{path}".')
-
-
-class InvalidCheckpointError(BrainPyError):
-    """A checkpoint cannot be stored in a directory that already has
-
-    a checkpoint at the current or a later step.
-
-    You can pass ``overwrite=True`` to disable this behavior and
-    overwrite existing checkpoints in the target directory.
-    """
-    __module__ = 'brainpy'
-
-    def __init__(self, path, step):
-        super().__init__(
-            f'Trying to save an outdated checkpoint at step: "{step}" and path: "{path}".'
-        )
-
-
-class InvalidCheckpointPath(BrainPyError):
-    """A checkpoint cannot be stored in a directory that already has
-
-    a checkpoint at the current or a later step.
-
-    You can pass ``overwrite=True`` to disable this behavior and
-    overwrite existing checkpoints in the target directory.
-    """
-    __module__ = 'brainpy'
-
-    def __init__(self, path):
-        super().__init__(f'Invalid checkpoint at "{path}".')
-
-
 class JaxTracerError(MathError):
     __module__ = 'brainpy'
 
@@ -271,22 +156,6 @@ def __init__(self, variables=None):
         super(JaxTracerError, self).__init__(msg)
 
 
-class ConcretizationTypeError(Exception):
-    __module__ = 'brainpy'
-
-    def __init__(self):
-        super(ConcretizationTypeError, self).__init__(
-            'This problem may be caused by several ways:\n'
-            '1. Your if-else conditional statement relies on instances of brainpy.math.Variable. \n'
-            '2. Your if-else conditional statement relies on functional arguments which do not '
-            'set in "static_argnames" when applying JIT compilation. More details please see '
-            'https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError\n'
-            '3. The static variables which set in the "static_argnames" are provided '
-            'as arguments, not keyword arguments, like "jit_f(v1, v2)" [<- wrong]. '
-            'Please write it as "jit_f(static_k1=v1, static_k2=v2)" [<- right].'
-        )
-
-
 class GPUOperatorNotFound(Exception):
     __module__ = 'brainpy'
 

brainpy/_exponential.py

@@ -1,4 +1,4 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2024 BrainX 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.
@@ -18,18 +18,20 @@
 
 from typing import Optional, Callable
 
-from brainstate.typing import Size, ArrayLike
-import brainstate
 import braintools
 import brainunit as u
+
+import brainstate
+from brainstate.typing import Size, ArrayLike
 from ._base import Synapse
+from .mixin import AlignPost
 
 __all__ = [
-     'Expon', 'DualExpon',
+    'Expon', 'DualExpon',
 ]
 
 
-class Expon(Synapse, brainstate.mixin.AlignPost):
+class Expon(Synapse, AlignPost):
     r"""
     Exponential decay synapse model.
 
@@ -97,7 +99,7 @@ def update(self, x=None):
         return self.g.value
 
 
-class DualExpon(Synapse, brainstate.mixin.AlignPost):
+class DualExpon(Synapse, AlignPost):
     r"""
     Dual exponential synapse model.
 

brainpy/_inputs.py

@@ -1,4 +1,4 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2024 BrainX 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.
@@ -21,9 +21,8 @@
 import numpy as np
 
 import brainstate
-from brainstate.typing import ArrayLike, Size, DTypeLike
 from brainpy._misc import set_module_as
-
+from brainstate.typing import ArrayLike, Size, DTypeLike
 
 __all__ = [
     'SpikeTime',

brainpy/_lif.py

@@ -1,4 +1,4 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2024 BrainX 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.
@@ -17,16 +17,16 @@
 
 from typing import Callable
 
-import brainstate
+import braintools
 import brainunit as u
 import jax
-from brainstate.typing import ArrayLike, Size
 
-import braintools
+import brainstate
+from brainstate.typing import ArrayLike, Size
 from ._base import Neuron
 
 __all__ = [
-     'IF', 'LIF', 'LIFRef', 'ALIF',
+    'IF', 'LIF', 'LIFRef', 'ALIF',
 ]
 
 

brainpy/_lif_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2024 BrainX 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.

brainpy/_misc.py

@@ -14,10 +14,9 @@
 # ==============================================================================
 
 
-
 def set_module_as(module: str):
     def wrapper(fun: callable):
         fun.__module__ = module
         return fun
 
-    return wrapper
+    return wrapper