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

brainstate/mixin.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.
@@ -15,35 +15,66 @@
 
 # -*- coding: utf-8 -*-
 
-from typing import (
-    Sequence, Optional, TypeVar, _SpecialForm, _type_check, _remove_dups_flatten, _UnionGenericAlias
-)
+"""
+Mixin classes and utility types for brainstate.
 
-import jax
+This module provides various mixin classes and custom type definitions that
+enhance the functionality of brainstate components. It includes parameter
+description mixins, alignment interfaces, and custom type definitions for
+expressing complex type requirements.
+"""
 
-T = TypeVar('T')
-ArrayLike = jax.typing.ArrayLike
+from typing import Sequence, Optional, TypeVar, Union, _GenericAlias
+
+import jax
 
 __all__ = [
     'Mixin',
     'ParamDesc',
     'ParamDescriber',
-    'AlignPost',
-    'BindCondData',
-
-    # types
     'JointTypes',
     'OneOfTypes',
-
-    # behavior modes
+    '_JointGenericAlias',
+    '_OneOfGenericAlias',
     'Mode',
     'JointMode',
     'Batching',
     'Training',
 ]
 
+T = TypeVar('T')
+ArrayLike = jax.typing.ArrayLike
+
 
 def hashable(x):
+    """
+    Check if an object is hashable.
+
+    Parameters
+    ----------
+    x : Any
+        The object to check for hashability.
+
+    Returns
+    -------
+    bool
+        True if the object is hashable, False otherwise.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate
+        >>>
+        >>> # Hashable objects
+        >>> assert brainstate.mixin.hashable(42) == True
+        >>> assert brainstate.mixin.hashable("string") == True
+        >>> assert brainstate.mixin.hashable((1, 2, 3)) == True
+        >>>
+        >>> # Non-hashable objects
+        >>> assert brainstate.mixin.hashable([1, 2, 3]) == False
+        >>> assert brainstate.mixin.hashable({"key": "value"}) == False
+    """
     try:
         hash(x)
         return True
@@ -52,45 +83,194 @@ def hashable(x):
 
 
 class Mixin(object):
-    """Base Mixin object.
-
-    The key for a :py:class:`~.Mixin` is that: no initialization function, only behavioral functions.
+    """
+    Base Mixin object for behavioral extensions.
+
+    The key characteristic of a :py:class:`~.Mixin` is that it provides only
+    behavioral functions without requiring initialization. Mixins are used to
+    add specific functionality to classes through multiple inheritance without
+    the complexity of a full base class.
+
+    Notes
+    -----
+    Mixins should not define ``__init__`` methods. They should only provide
+    methods that add specific behaviors to the classes that inherit from them.
+
+    Examples
+    --------
+    Creating a custom mixin:
+
+    .. code-block:: python
+
+        >>> import brainstate
+        >>>
+        >>> class LoggingMixin(brainstate.mixin.Mixin):
+        ...     def log(self, message):
+        ...         print(f"[{self.__class__.__name__}] {message}")
+
+        >>> class MyComponent(brainstate.nn.Module, LoggingMixin):
+        ...     def __init__(self):
+        ...         super().__init__()
+        ...
+        ...     def process(self):
+        ...         self.log("Processing data...")
+        ...         return "Done"
+        >>>
+        >>> component = MyComponent()
+        >>> component.process()  # Prints: [MyComponent] Processing data...
     """
     pass
 
 
 class ParamDesc(Mixin):
     """
-    :py:class:`~.Mixin` indicates the function for describing initialization parameters.
-
-    This mixin enables the subclass has a classmethod ``desc``, which
-    produces an instance of :py:class:`~.ParamDescriber`.
-
-    Note this Mixin can be applied in any Python object.
+    Mixin for describing initialization parameters.
+
+    This mixin enables a class to have a ``desc`` classmethod, which produces
+    an instance of :py:class:`~.ParamDescriber`. This is useful for creating
+    parameter templates that can be reused to instantiate multiple objects
+    with the same configuration.
+
+    Attributes
+    ----------
+    non_hashable_params : sequence of str, optional
+        Names of parameters that are not hashable and should be handled specially.
+
+    Notes
+    -----
+    This mixin can be applied to any Python class, not just brainstate-specific classes.
+
+    Examples
+    --------
+    Basic usage of ParamDesc:
+
+    .. code-block:: python
+
+        >>> import brainstate
+        >>>
+        >>> class NeuronModel(brainstate.mixin.ParamDesc):
+        ...     def __init__(self, size, tau=10.0, threshold=1.0):
+        ...         self.size = size
+        ...         self.tau = tau
+        ...         self.threshold = threshold
+        >>>
+        >>> # Create a parameter descriptor
+        >>> neuron_desc = NeuronModel.desc(size=100, tau=20.0)
+        >>>
+        >>> # Use the descriptor to create instances
+        >>> neuron1 = neuron_desc(threshold=0.8)  # Creates with threshold=0.8
+        >>> neuron2 = neuron_desc(threshold=1.2)  # Creates with threshold=1.2
+        >>>
+        >>> # Both neurons share size=100, tau=20.0 but have different thresholds
+
+    Creating reusable templates:
+
+    .. code-block:: python
+
+        >>> # Define a template for excitatory neurons
+        >>> exc_neuron_template = NeuronModel.desc(size=1000, tau=10.0, threshold=1.0)
+        >>>
+        >>> # Define a template for inhibitory neurons
+        >>> inh_neuron_template = NeuronModel.desc(size=250, tau=5.0, threshold=0.5)
+        >>>
+        >>> # Create multiple instances from templates
+        >>> exc_population = [exc_neuron_template() for _ in range(5)]
+        >>> inh_population = [inh_neuron_template() for _ in range(2)]
     """
 
+    # Optional list of parameter names that are not hashable
+    # These will be converted to strings for hashing purposes
     non_hashable_params: Optional[Sequence[str]] = None
 
     @classmethod
     def desc(cls, *args, **kwargs) -> 'ParamDescriber':
+        """
+        Create a parameter describer for this class.
+
+        Parameters
+        ----------
+        *args
+            Positional arguments to be used in future instantiations.
+        **kwargs
+            Keyword arguments to be used in future instantiations.
+
+        Returns
+        -------
+        ParamDescriber
+            A descriptor that can be used to create instances with these parameters.
+        """
         return ParamDescriber(cls, *args, **kwargs)
 
 
 class HashableDict(dict):
+    """
+    A dictionary that can be hashed by converting non-hashable values to strings.
+
+    This is used internally to make parameter dictionaries hashable so they can
+    be used as part of cache keys or other contexts requiring hashability.
+
+    Parameters
+    ----------
+    the_dict : dict
+        The dictionary to make hashable.
+
+    Notes
+    -----
+    Non-hashable values in the dictionary are automatically converted to their
+    string representation.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate
+        >>> import jax.numpy as jnp
+        >>>
+        >>> # Regular dict with non-hashable values cannot be hashed
+        >>> regular_dict = {"array": jnp.array([1, 2, 3]), "value": 42}
+        >>> # hash(regular_dict)  # This would raise TypeError
+        >>>
+        >>> # HashableDict can be hashed
+        >>> hashable = brainstate.mixin.HashableDict(regular_dict)
+        >>> key = hash(hashable)  # This works!
+        >>>
+        >>> # Can be used in sets or as dict keys
+        >>> cache = {hashable: "result"}
+    """
+
     def __init__(self, the_dict: dict):
+        # Process the dictionary to ensure all values are hashable
         out = dict()
         for k, v in the_dict.items():
             if not hashable(v):
-                v = str(v)  # convert to string if not hashable
+                # Convert non-hashable values to their string representation
+                v = str(v)
             out[k] = v
         super().__init__(out)
 
     def __hash__(self):
+        """
+        Compute hash from sorted items for consistent hashing regardless of insertion order.
+        """
         return hash(tuple(sorted(self.items())))
 
 
 class NoSubclassMeta(type):
+    """
+    Metaclass that prevents a class from being subclassed.
+
+    This is used to ensure that certain classes (like ParamDescriber) are used
+    as-is and not extended through inheritance, which could lead to unexpected
+    behavior.
+
+    Raises
+    ------
+    TypeError
+        If an attempt is made to subclass a class using this metaclass.
+    """
+
     def __new__(cls, name, bases, classdict):
+        # Check if any base class uses NoSubclassMeta
         for b in bases:
             if isinstance(b, NoSubclassMeta):
                 raise TypeError("type '{0}' is not an acceptable base type".format(b.__name__))
@@ -99,209 +279,758 @@ class NoSubclassMeta(type):
 
 class ParamDescriber(metaclass=NoSubclassMeta):
     """
-    ParamDesc initialization for parameter describers.
+    Parameter descriptor for deferred object instantiation.
+
+    This class stores a class reference along with arguments and keyword arguments,
+    allowing for deferred instantiation. It's useful for creating templates that
+    can be reused to create multiple instances with similar configurations.
+
+    Parameters
+    ----------
+    cls : type
+        The class to be instantiated.
+    *desc_tuple
+        Positional arguments to be stored and used during instantiation.
+    **desc_dict
+        Keyword arguments to be stored and used during instantiation.
+
+    Attributes
+    ----------
+    cls : type
+        The class that will be instantiated.
+    args : tuple
+        Stored positional arguments.
+    kwargs : dict
+        Stored keyword arguments.
+    identifier : tuple
+        A hashable identifier for this descriptor.
+
+    Notes
+    -----
+    ParamDescriber cannot be subclassed due to the NoSubclassMeta metaclass.
+    This ensures consistent behavior across the codebase.
+
+    Examples
+    --------
+    Manual creation of a descriptor:
+
+    .. code-block:: python
+
+        >>> import brainstate
+        >>>
+        >>> class Network:
+        ...     def __init__(self, n_neurons, learning_rate=0.01):
+        ...         self.n_neurons = n_neurons
+        ...         self.learning_rate = learning_rate
+        >>>
+        >>> # Create a descriptor
+        >>> network_desc = brainstate.mixin.ParamDescriber(
+        ...     Network, n_neurons=1000, learning_rate=0.001
+        ... )
+        >>>
+        >>> # Use the descriptor to create instances with additional args
+        >>> net1 = network_desc()
+        >>> net2 = network_desc()  # Same configuration
+
+    Using with ParamDesc mixin:
+
+    .. code-block:: python
+
+        >>> class Network(brainstate.mixin.ParamDesc):
+        ...     def __init__(self, n_neurons, learning_rate=0.01):
+        ...         self.n_neurons = n_neurons
+        ...         self.learning_rate = learning_rate
+        >>>
+        >>> # More concise syntax using the desc() classmethod
+        >>> network_desc = Network.desc(n_neurons=1000)
+        >>> net = network_desc(learning_rate=0.005)  # Override learning_rate
     """
 
     def __init__(self, cls: T, *desc_tuple, **desc_dict):
+        # Store the class to be instantiated
         self.cls: type = cls
 
-        # arguments
+        # Store the arguments for later instantiation
         self.args = desc_tuple
         self.kwargs = desc_dict
 
-        # identifier
+        # Create a hashable identifier for caching/comparison purposes
+        # This combines the class, args tuple, and hashable kwargs dict
         self._identifier = (cls, tuple(desc_tuple), HashableDict(desc_dict))
 
     def __call__(self, *args, **kwargs) -> T:
-        return self.cls(*self.args, *args, **self.kwargs, **kwargs)
+        """
+        Instantiate the class with stored and additional arguments.
+
+        Parameters
+        ----------
+        *args
+            Additional positional arguments to append.
+        **kwargs
+            Additional keyword arguments to merge (will override stored kwargs).
+
+        Returns
+        -------
+        T
+            An instance of the described class.
+        """
+        # Merge stored arguments with new arguments
+        # Stored args come first, then new args
+        # Merge kwargs with new kwargs overriding stored ones
+        merged_kwargs = {**self.kwargs, **kwargs}
+        return self.cls(*self.args, *args, **merged_kwargs)
 
     def init(self, *args, **kwargs):
+        """
+        Alias for __call__, explicitly named for clarity.
+
+        Parameters
+        ----------
+        *args
+            Additional positional arguments.
+        **kwargs
+            Additional keyword arguments.
+
+        Returns
+        -------
+        T
+            An instance of the described class.
+        """
         return self.__call__(*args, **kwargs)
 
     def __instancecheck__(self, instance):
+        """
+        Check if an instance is compatible with this descriptor.
+
+        Parameters
+        ----------
+        instance : Any
+            The instance to check.
+
+        Returns
+        -------
+        bool
+            True if the instance is a ParamDescriber for a compatible class.
+        """
+        # Must be a ParamDescriber
         if not isinstance(instance, ParamDescriber):
             return False
+        # The described class must be a subclass of our class
         if not issubclass(instance.cls, self.cls):
             return False
         return True
 
     @classmethod
     def __class_getitem__(cls, item: type):
+        """
+        Support for subscript notation: ParamDescriber[MyClass].
+
+        Parameters
+        ----------
+        item : type
+            The class to create a descriptor for.
+
+        Returns
+        -------
+        ParamDescriber
+            A descriptor for the given class.
+        """
         return ParamDescriber(item)
 
     @property
     def identifier(self):
+        """
+        Get the unique identifier for this descriptor.
+
+        Returns
+        -------
+        tuple
+            A hashable identifier consisting of (class, args, kwargs).
+        """
         return self._identifier
 
     @identifier.setter
     def identifier(self, value: ArrayLike):
+        """
+        Prevent modification of the identifier.
+
+        Raises
+        ------
+        AttributeError
+            Always, as the identifier is read-only.
+        """
         raise AttributeError('Cannot set the identifier.')
 
 
-class AlignPost(Mixin):
+def not_implemented(func):
     """
-    Align post MixIn.
-
-    This class provides a ``align_post_input_add()`` function for
-    add external currents.
+    Decorator to mark a function as not implemented.
+
+    This decorator wraps a function to raise NotImplementedError when called,
+    and adds a ``not_implemented`` attribute for checking.
+
+    Parameters
+    ----------
+    func : callable
+        The function to mark as not implemented.
+
+    Returns
+    -------
+    callable
+        A wrapper function that raises NotImplementedError.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate
+        >>>
+        >>> class BaseModel:
+        ...     @brainstate.mixin.not_implemented
+        ...     def process(self, x):
+        ...         pass
+        >>>
+        >>> model = BaseModel()
+        >>> # model.process(10)  # Raises: NotImplementedError: process is not implemented.
+        >>>
+        >>> # Check if a method is not implemented
+        >>> assert hasattr(BaseModel.process, 'not_implemented')
     """
 
-    def align_post_input_add(self, *args, **kwargs):
-        raise NotImplementedError
-
-
-class BindCondData(Mixin):
-    """Bind temporary conductance data.
-
-
-    """
-    _conductance: Optional
-
-    def bind_cond(self, conductance):
-        self._conductance = conductance
-
-    def unbind_cond(self):
-        self._conductance = None
-
-
-def not_implemented(func):
     def wrapper(*args, **kwargs):
         raise NotImplementedError(f'{func.__name__} is not implemented.')
 
+    # Mark the wrapper so we can detect not-implemented methods
     wrapper.not_implemented = True
     return wrapper
 
 
-class _MetaUnionType(type):
-    def __new__(cls, name, bases, dct):
-        if isinstance(bases, type):
-            bases = (bases,)
-        elif isinstance(bases, (list, tuple)):
-            bases = tuple(bases)
-            for base in bases:
-                assert isinstance(base, type), f'Must be type. But got {base}'
-        else:
-            raise TypeError(f'Must be type. But got {bases}')
-        return super().__new__(cls, name, bases, dct)
+class _JointGenericAlias(_GenericAlias, _root=True):
+    """
+    Generic alias for JointTypes (intersection types).
+
+    This class represents a type that requires all specified types to be satisfied.
+    Unlike _MetaUnionType which creates actual classes with metaclass conflicts,
+    this uses typing's generic alias system to avoid metaclass issues.
+    """
 
-    def __instancecheck__(self, other):
-        cls_of_other = other.__class__
-        return all([issubclass(cls_of_other, cls) for cls in self.__bases__])
+    def __instancecheck__(self, obj):
+        """
+        Check if an instance is an instance of all component types.
+        """
+        return all(isinstance(obj, cls) for cls in self.__args__)
 
     def __subclasscheck__(self, subclass):
-        return all([issubclass(subclass, cls) for cls in self.__bases__])
+        """
+        Check if a class is a subclass of all component types.
+        """
+        return all(issubclass(subclass, cls) for cls in self.__args__)
+
+    def __eq__(self, other):
+        """
+        Check equality with another type.
 
+        Two JointTypes are equal if they have the same component types,
+        regardless of order.
+        """
+        if not isinstance(other, _JointGenericAlias):
+            return NotImplemented
+        return set(self.__args__) == set(other.__args__)
 
-class _JointGenericAlias(_UnionGenericAlias, _root=True):
-    def __subclasscheck__(self, subclass):
-        return all([issubclass(subclass, cls) for cls in set(self.__args__)])
+    def __hash__(self):
+        """
+        Return hash of the JointType.
 
+        The hash is based on the frozenset of component types to ensure
+        that JointTypes with the same types (regardless of order) have
+        the same hash.
+        """
+        return hash(frozenset(self.__args__))
 
-@_SpecialForm
-def JointTypes(self, parameters):
-    """Joint types; JointTypes[X, Y] means both X and Y.
+    def __repr__(self):
+        """
+        Return string representation of the JointType.
 
-    To define a union, use e.g. Union[int, str].
+        Returns a readable representation showing all component types.
+        """
+        args_str = ', '.join(
+            arg.__module__ + '.' + arg.__name__ if hasattr(arg, '__module__') and hasattr(arg, '__name__')
+            else str(arg)
+            for arg in self.__args__
+        )
+        return f'JointTypes[{args_str}]'
+
+    def __reduce__(self):
+        """
+        Support for pickling.
 
-    Details:
-    - The arguments must be types and there must be at least one.
-    - None as an argument is a special case and is replaced by `type(None)`.
-    - Unions of unions are flattened, e.g.::
+        Returns the necessary information to reconstruct the JointType
+        when unpickling.
+        """
+        return (_JointGenericAlias, (self.__origin__, self.__args__))
 
-        JointTypes[JointTypes[int, str], float] == JointTypes[int, str, float]
 
-    - Unions of a single argument vanish, e.g.::
+class _OneOfGenericAlias(_GenericAlias, _root=True):
+    """
+    Generic alias for OneOfTypes (union types).
 
-        JointTypes[int] == int  # The constructor actually returns int
+    This class represents a type that requires at least one of the specified
+    types to be satisfied. It's similar to typing.Union but provides a consistent
+    interface with JointTypes and avoids potential metaclass conflicts.
+    """
 
-    - Redundant arguments are skipped, e.g.::
+    def __instancecheck__(self, obj):
+        """
+        Check if an instance is an instance of any component type.
+        """
+        return any(isinstance(obj, cls) for cls in self.__args__)
 
-        JointTypes[int, str, int] == JointTypes[int, str]
+    def __subclasscheck__(self, subclass):
+        """
+        Check if a class is a subclass of any component type.
+        """
+        return any(issubclass(subclass, cls) for cls in self.__args__)
 
-    - When comparing unions, the argument order is ignored, e.g.::
+    def __eq__(self, other):
+        """
+        Check equality with another type.
 
-        JointTypes[int, str] == JointTypes[str, int]
+        Two OneOfTypes are equal if they have the same component types,
+        regardless of order.
+        """
+        if not isinstance(other, _OneOfGenericAlias):
+            return NotImplemented
+        return set(self.__args__) == set(other.__args__)
 
-    - You cannot subclass or instantiate a JointTypes.
-    - You can use Optional[X] as a shorthand for JointTypes[X, None].
-    """
-    if parameters == ():
-        raise TypeError("Cannot take a Joint of no types.")
-    if not isinstance(parameters, tuple):
-        parameters = (parameters,)
-    msg = "JointTypes[arg, ...]: each arg must be a type."
-    parameters = tuple(_type_check(p, msg) for p in parameters)
-    parameters = _remove_dups_flatten(parameters)
-    if len(parameters) == 1:
-        return parameters[0]
-    if len(parameters) == 2 and type(None) in parameters:
-        return _UnionGenericAlias(self, parameters, name="Optional")
-    return _JointGenericAlias(self, parameters)
+    def __hash__(self):
+        """
+        Return hash of the OneOfType.
 
+        The hash is based on the frozenset of component types to ensure
+        that OneOfTypes with the same types (regardless of order) have
+        the same hash.
+        """
+        return hash(frozenset(self.__args__))
 
-@_SpecialForm
-def OneOfTypes(self, parameters):
-    """Sole type; OneOfTypes[X, Y] means either X or Y.
+    def __repr__(self):
+        """
+        Return string representation of the OneOfType.
 
-    To define a union, use e.g. OneOfTypes[int, str]. Details:
-    - The arguments must be types and there must be at least one.
-    - None as an argument is a special case and is replaced by
-      type(None).
-    - Unions of unions are flattened, e.g.::
+        Returns a readable representation showing all component types.
+        """
+        args_str = ', '.join(
+            arg.__module__ + '.' + arg.__name__ if hasattr(arg, '__module__') and hasattr(arg, '__name__')
+            else str(arg)
+            for arg in self.__args__
+        )
+        return f'OneOfTypes[{args_str}]'
+
+    def __reduce__(self):
+        """
+        Support for pickling.
 
-        assert OneOfTypes[OneOfTypes[int, str], float] == OneOfTypes[int, str, float]
+        Returns the necessary information to reconstruct the OneOfType
+        when unpickling.
+        """
+        return (_OneOfGenericAlias, (self.__origin__, self.__args__))
 
-    - Unions of a single argument vanish, e.g.::
 
-        assert OneOfTypes[int] == int  # The constructor actually returns int
+class _JointTypesClass:
+    """Helper class to enable subscript syntax for JointTypes."""
 
-    - Redundant arguments are skipped, e.g.::
+    def __call__(self, *types):
+        """
+        Create a type that requires all specified types (intersection type).
+
+        This function creates a type hint that indicates a value must satisfy all
+        the specified types simultaneously. It's useful for expressing complex
+        type requirements where a single object must implement multiple interfaces.
+
+        Parameters
+        ----------
+        *types : type
+            The types that must all be satisfied.
+
+        Returns
+        -------
+        type
+            A type that checks for all specified types.
+
+        Notes
+        -----
+        - If only one type is provided, that type is returned directly.
+        - Redundant types are automatically removed.
+        - The order of types doesn't matter for equality checks.
+
+        Examples
+        --------
+        Basic usage with interfaces:
+
+        .. code-block:: python
+
+            >>> import brainstate
+            >>> from typing import Protocol
+            >>>
+            >>> class Trainable(Protocol):
+            ...     def train(self): ...
+            >>>
+            >>> class Evaluable(Protocol):
+            ...     def evaluate(self): ...
+            >>>
+            >>> # A model that is both trainable and evaluable
+            >>> TrainableEvaluableModel = brainstate.mixin.JointTypes(Trainable, Evaluable)
+            >>> # Or using subscript syntax
+            >>> TrainableEvaluableModel = brainstate.mixin.JointTypes[Trainable, Evaluable]
+            >>>
+            >>> class NeuralNetwork(Trainable, Evaluable):
+            ...     def train(self):
+            ...         return "Training..."
+            ...
+            ...     def evaluate(self):
+            ...         return "Evaluating..."
+            >>>
+            >>> model = NeuralNetwork()
+            >>> # model satisfies JointTypes(Trainable, Evaluable)
+
+        Using with mixin classes:
+
+        .. code-block:: python
+
+            >>> class Serializable:
+            ...     def save(self): pass
+            >>>
+            >>> class Visualizable:
+            ...     def plot(self): pass
+            >>>
+            >>> # Require both serialization and visualization
+            >>> FullFeaturedModel = brainstate.mixin.JointTypes[Serializable, Visualizable]
+            >>>
+            >>> class MyModel(Serializable, Visualizable):
+            ...     def save(self):
+            ...         return "Saved"
+            ...
+            ...     def plot(self):
+            ...         return "Plotted"
+        """
+        if len(types) == 0:
+            raise TypeError("Cannot create a JointTypes of no types.")
+
+        # Remove duplicates while preserving some order
+        seen = set()
+        unique_types = []
+        for t in types:
+            if t not in seen:
+                seen.add(t)
+                unique_types.append(t)
+
+        # If only one type, return it directly
+        if len(unique_types) == 1:
+            return unique_types[0]
+
+        # Create a generic alias for the joint type
+        # This avoids metaclass conflicts by using typing's generic alias system
+        return _JointGenericAlias(object, tuple(unique_types))
+
+    def __getitem__(self, item):
+        """Enable subscript syntax: JointTypes[Type1, Type2]."""
+        if isinstance(item, tuple):
+            return self(*item)
+        else:
+            return self(item)
 
-        assert OneOfTypes[int, str, int] == OneOfTypes[int, str]
 
-    - When comparing unions, the argument order is ignored, e.g.::
+# Create singleton instance that acts as both a callable and supports subscript
+JointTypes = _JointTypesClass()
 
-        assert OneOfTypes[int, str] == OneOfTypes[str, int]
 
-    - You cannot subclass or instantiate a union.
-    - You can use Optional[X] as a shorthand for OneOfTypes[X, None].
-    """
-    if parameters == ():
-        raise TypeError("Cannot take a Sole of no types.")
-    if not isinstance(parameters, tuple):
-        parameters = (parameters,)
-    msg = "OneOfTypes[arg, ...]: each arg must be a type."
-    parameters = tuple(_type_check(p, msg) for p in parameters)
-    parameters = _remove_dups_flatten(parameters)
-    if len(parameters) == 1:
-        return parameters[0]
-    if len(parameters) == 2 and type(None) in parameters:
-        return _UnionGenericAlias(self, parameters, name="Optional")
-    return _UnionGenericAlias(self, parameters)
+class _OneOfTypesClass:
+    """Helper class to enable subscript syntax for OneOfTypes."""
+
+    def __call__(self, *types):
+        """
+        Create a type that requires one of the specified types (union type).
+
+        This is similar to typing.Union but provides a more intuitive name and
+        consistent behavior with JointTypes. It indicates that a value must satisfy
+        at least one of the specified types.
+
+        Parameters
+        ----------
+        *types : type
+            The types, one of which must be satisfied.
+
+        Returns
+        -------
+        Union type
+            A union type of the specified types.
+
+        Notes
+        -----
+        - If only one type is provided, that type is returned directly.
+        - Redundant types are automatically removed.
+        - The order of types doesn't matter for equality checks.
+        - This is equivalent to typing.Union[...].
+
+        Examples
+        --------
+        Basic usage with different types:
+
+        .. code-block:: python
+
+            >>> import brainstate
+            >>>
+            >>> # A parameter that can be int or float
+            >>> NumericType = brainstate.mixin.OneOfTypes(int, float)
+            >>> # Or using subscript syntax
+            >>> NumericType = brainstate.mixin.OneOfTypes[int, float]
+            >>>
+            >>> def process_value(x: NumericType):
+            ...     return x * 2
+            >>>
+            >>> # Both work
+            >>> result1 = process_value(5)      # int
+            >>> result2 = process_value(3.14)   # float
+
+        Using with class types:
+
+        .. code-block:: python
+
+            >>> class NumpyArray:
+            ...     pass
+            >>>
+            >>> class JAXArray:
+            ...     pass
+            >>>
+            >>> # Accept either numpy or JAX arrays
+            >>> ArrayType = brainstate.mixin.OneOfTypes[NumpyArray, JAXArray]
+            >>>
+            >>> def compute(arr: ArrayType):
+            ...     if isinstance(arr, NumpyArray):
+            ...         return "Processing numpy array"
+            ...     elif isinstance(arr, JAXArray):
+            ...         return "Processing JAX array"
+
+        Combining with None for optional types:
+
+        .. code-block:: python
+
+            >>> # Optional string (equivalent to Optional[str])
+            >>> MaybeString = brainstate.mixin.OneOfTypes[str, type(None)]
+            >>>
+            >>> def format_name(name: MaybeString) -> str:
+            ...     if name is None:
+            ...         return "Anonymous"
+            ...     return name.title()
+        """
+        if len(types) == 0:
+            raise TypeError("Cannot create a OneOfTypes of no types.")
+
+        # Remove duplicates
+        seen = set()
+        unique_types = []
+        for t in types:
+            if t not in seen:
+                seen.add(t)
+                unique_types.append(t)
+
+        # If only one type, return it directly
+        if len(unique_types) == 1:
+            return unique_types[0]
+
+        # Create a generic alias for the union type
+        # This provides consistency with JointTypes and avoids metaclass conflicts
+        return _OneOfGenericAlias(Union, tuple(unique_types))
+
+    def __getitem__(self, item):
+        """Enable subscript syntax: OneOfTypes[Type1, Type2]."""
+        if isinstance(item, tuple):
+            return self(*item)
+        else:
+            return self(item)
+
+
+# Create singleton instance that acts as both a callable and supports subscript
+OneOfTypes = _OneOfTypesClass()
+
+
+def __getattr__(name):
+    if name in [
+        'Mode',
+        'JointMode',
+        'Batching',
+        'Training',
+        'AlignPost',
+        'BindCondData',
+    ]:
+        import warnings
+        warnings.warn(
+            f"brainstate.mixin.{name} is deprecated and will be removed in a future version. "
+            f"Please use brainpy.mixin.{name} instead.",
+            DeprecationWarning,
+            stacklevel=2
+        )
+        import brainpy
+        return getattr(brainpy.mixin, name)
+    raise AttributeError(
+        f'module {__name__!r} has no attribute {name!r}'
+    )
 
 
 class Mode(Mixin):
     """
-    Base class for computation behaviors.
+    Base class for computation behavior modes.
+
+    Modes are used to represent different computational contexts or behaviors,
+    such as training vs evaluation, batched vs single-sample processing, etc.
+    They provide a flexible way to configure how models and components behave
+    in different scenarios.
+
+    Examples
+    --------
+    Creating a custom mode:
+
+    .. code-block:: python
+
+        >>> import brainstate
+        >>>
+        >>> class InferenceMode(brainstate.mixin.Mode):
+        ...     def __init__(self, use_cache=True):
+        ...         self.use_cache = use_cache
+        >>>
+        >>> # Create mode instances
+        >>> inference = InferenceMode(use_cache=True)
+        >>> print(inference)  # Output: InferenceMode
+
+    Checking mode types:
+
+    .. code-block:: python
+
+        >>> class FastMode(brainstate.mixin.Mode):
+        ...     pass
+        >>>
+        >>> class SlowMode(brainstate.mixin.Mode):
+        ...     pass
+        >>>
+        >>> fast = FastMode()
+        >>> slow = SlowMode()
+        >>>
+        >>> # Check exact mode type
+        >>> assert fast.is_a(FastMode)
+        >>> assert not fast.is_a(SlowMode)
+        >>>
+        >>> # Check if mode is an instance of a type
+        >>> assert fast.has(brainstate.mixin.Mode)
+
+    Using modes in a model:
+
+    .. code-block:: python
+
+        >>> class Model:
+        ...     def __init__(self):
+        ...         self.mode = brainstate.mixin.Training()
+        ...
+        ...     def forward(self, x):
+        ...         if self.mode.has(brainstate.mixin.Training):
+        ...             # Training-specific logic
+        ...             return self.train_forward(x)
+        ...         else:
+        ...             # Inference logic
+        ...             return self.eval_forward(x)
+        ...
+        ...     def train_forward(self, x):
+        ...         return x + 0.1  # Add noise during training
+        ...
+        ...     def eval_forward(self, x):
+        ...         return x  # No noise during evaluation
     """
 
     def __repr__(self):
+        """
+        String representation of the mode.
+
+        Returns
+        -------
+        str
+            The class name of the mode.
+        """
         return self.__class__.__name__
 
     def __eq__(self, other: 'Mode'):
+        """
+        Check equality of modes based on their type.
+
+        Parameters
+        ----------
+        other : Mode
+            Another mode to compare with.
+
+        Returns
+        -------
+        bool
+            True if both modes are of the same class.
+        """
         assert isinstance(other, Mode)
         return other.__class__ == self.__class__
 
     def is_a(self, mode: type):
         """
-        Check whether the mode is exactly the desired mode.
+        Check whether the mode is exactly the desired mode type.
+
+        This performs an exact type match, not checking for subclasses.
+
+        Parameters
+        ----------
+        mode : type
+            The mode type to check against.
+
+        Returns
+        -------
+        bool
+            True if this mode is exactly of the specified type.
+
+        Examples
+        --------
+        .. code-block:: python
+
+            >>> import brainstate
+            >>>
+            >>> training_mode = brainstate.mixin.Training()
+            >>> assert training_mode.is_a(brainstate.mixin.Training)
+            >>> assert not training_mode.is_a(brainstate.mixin.Batching)
         """
         assert isinstance(mode, type), 'Must be a type.'
         return self.__class__ == mode
 
     def has(self, mode: type):
         """
-        Check whether the mode is included in the desired mode.
+        Check whether the mode includes the desired mode type.
+
+        This checks if the current mode is an instance of the specified type,
+        including checking for subclasses.
+
+        Parameters
+        ----------
+        mode : type
+            The mode type to check for.
+
+        Returns
+        -------
+        bool
+            True if this mode is an instance of the specified type.
+
+        Examples
+        --------
+        .. code-block:: python
+
+            >>> import brainstate
+            >>>
+            >>> # Create a custom mode that extends Training
+            >>> class AdvancedTraining(brainstate.mixin.Training):
+            ...     pass
+            >>>
+            >>> advanced = AdvancedTraining()
+            >>> assert advanced.has(brainstate.mixin.Training)  # True (subclass)
+            >>> assert advanced.has(brainstate.mixin.Mode)      # True (base class)
         """
         assert isinstance(mode, type), 'Must be a type.'
         return isinstance(self, mode)
@@ -309,57 +1038,396 @@ class Mode(Mixin):
 
 class JointMode(Mode):
     """
-    Joint mode.
+    A mode that combines multiple modes simultaneously.
+
+    JointMode allows expressing that a computation is in multiple modes at once,
+    such as being both in training mode and batching mode. This is useful for
+    complex scenarios where multiple behavioral aspects need to be active.
+
+    Parameters
+    ----------
+    *modes : Mode
+        The modes to combine.
+
+    Attributes
+    ----------
+    modes : tuple of Mode
+        The individual modes that are combined.
+    types : set of type
+        The types of the combined modes.
+
+    Raises
+    ------
+    TypeError
+        If any of the provided arguments is not a Mode instance.
+
+    Examples
+    --------
+    Combining training and batching modes:
+
+    .. code-block:: python
+
+        >>> import brainstate
+        >>>
+        >>> # Create individual modes
+        >>> training = brainstate.mixin.Training()
+        >>> batching = brainstate.mixin.Batching(batch_size=32)
+        >>>
+        >>> # Combine them
+        >>> joint = brainstate.mixin.JointMode(training, batching)
+        >>> print(joint)  # JointMode(Training, Batching(in_size=32, axis=0))
+        >>>
+        >>> # Check if specific modes are present
+        >>> assert joint.has(brainstate.mixin.Training)
+        >>> assert joint.has(brainstate.mixin.Batching)
+        >>>
+        >>> # Access attributes from combined modes
+        >>> print(joint.batch_size)  # 32 (from Batching mode)
+
+    Using in model configuration:
+
+    .. code-block:: python
+
+        >>> class NeuralNetwork:
+        ...     def __init__(self):
+        ...         self.mode = None
+        ...
+        ...     def set_train_mode(self, batch_size=1):
+        ...         # Set both training and batching modes
+        ...         training = brainstate.mixin.Training()
+        ...         batching = brainstate.mixin.Batching(batch_size=batch_size)
+        ...         self.mode = brainstate.mixin.JointMode(training, batching)
+        ...
+        ...     def forward(self, x):
+        ...         if self.mode.has(brainstate.mixin.Training):
+        ...             x = self.apply_dropout(x)
+        ...
+        ...         if self.mode.has(brainstate.mixin.Batching):
+        ...             # Process in batches
+        ...             batch_size = self.mode.batch_size
+        ...             return self.batch_process(x, batch_size)
+        ...
+        ...         return self.process(x)
+        >>>
+        >>> model = NeuralNetwork()
+        >>> model.set_train_mode(batch_size=64)
     """
 
     def __init__(self, *modes: Mode):
+        # Validate that all arguments are Mode instances
         for m_ in modes:
             if not isinstance(m_, Mode):
                 raise TypeError(f'The supported type must be a tuple/list of Mode. But we got {m_}')
+
+        # Store the modes as a tuple
         self.modes = tuple(modes)
+
+        # Store the types of the modes for quick lookup
         self.types = set([m.__class__ for m in modes])
 
     def __repr__(self):
+        """
+        String representation showing all combined modes.
+
+        Returns
+        -------
+        str
+            A string showing the joint mode and its components.
+        """
         return f'{self.__class__.__name__}({", ".join([repr(m) for m in self.modes])})'
 
     def has(self, mode: type):
         """
-        Check whether the mode is included in the desired mode.
+        Check whether any of the combined modes includes the desired type.
+
+        Parameters
+        ----------
+        mode : type
+            The mode type to check for.
+
+        Returns
+        -------
+        bool
+            True if any of the combined modes is or inherits from the specified type.
+
+        Examples
+        --------
+        .. code-block:: python
+
+            >>> import brainstate
+            >>>
+            >>> training = brainstate.mixin.Training()
+            >>> batching = brainstate.mixin.Batching(batch_size=16)
+            >>> joint = brainstate.mixin.JointMode(training, batching)
+            >>>
+            >>> assert joint.has(brainstate.mixin.Training)
+            >>> assert joint.has(brainstate.mixin.Batching)
+            >>> assert joint.has(brainstate.mixin.Mode)  # Base class
         """
         assert isinstance(mode, type), 'Must be a type.'
+        # Check if any of the combined mode types is a subclass of the target mode
         return any([issubclass(cls, mode) for cls in self.types])
 
     def is_a(self, cls: type):
         """
-        Check whether the mode is exactly the desired mode.
+        Check whether the joint mode is exactly the desired combined type.
+
+        This is a complex check that verifies the joint mode matches a specific
+        combination of types.
+
+        Parameters
+        ----------
+        cls : type
+            The combined type to check against.
+
+        Returns
+        -------
+        bool
+            True if the joint mode exactly matches the specified type combination.
         """
-        return JointTypes[tuple(self.types)] == cls
+        # Use JointTypes to create the expected type from our mode types
+        return JointTypes(*tuple(self.types)) == cls
 
     def __getattr__(self, item):
         """
-        Get the attribute from the mode.
-
-        If the attribute is not found in the mode, then it will be searched in the base class.
+        Get attributes from the combined modes.
+
+        This method searches through all combined modes to find the requested
+        attribute, allowing transparent access to properties of any of the
+        combined modes.
+
+        Parameters
+        ----------
+        item : str
+            The attribute name to search for.
+
+        Returns
+        -------
+        Any
+            The attribute value from the first mode that has it.
+
+        Raises
+        ------
+        AttributeError
+            If the attribute is not found in any of the combined modes.
+
+        Examples
+        --------
+        .. code-block:: python
+
+            >>> import brainstate
+            >>>
+            >>> batching = brainstate.mixin.Batching(batch_size=32, batch_axis=1)
+            >>> training = brainstate.mixin.Training()
+            >>> joint = brainstate.mixin.JointMode(batching, training)
+            >>>
+            >>> # Access batching attributes directly
+            >>> print(joint.batch_size)  # 32
+            >>> print(joint.batch_axis)  # 1
         """
+        # Don't interfere with accessing modes and types attributes
         if item in ['modes', 'types']:
             return super().__getattribute__(item)
+
+        # Search for the attribute in each combined mode
         for m in self.modes:
             if hasattr(m, item):
                 return getattr(m, item)
+
+        # If not found, fall back to default behavior (will raise AttributeError)
         return super().__getattribute__(item)
 
 
 class Batching(Mode):
-    """Batching mode."""
+    """
+    Mode indicating batched computation.
+
+    This mode specifies that computations should be performed on batches of data,
+    including information about the batch size and which axis represents the batch
+    dimension.
+
+    Parameters
+    ----------
+    batch_size : int, default 1
+        The size of each batch.
+    batch_axis : int, default 0
+        The axis along which batching occurs.
+
+    Attributes
+    ----------
+    batch_size : int
+        The number of samples in each batch.
+    batch_axis : int
+        The axis index representing the batch dimension.
+
+    Examples
+    --------
+    Basic batching configuration:
+
+    .. code-block:: python
+
+        >>> import brainstate
+        >>>
+        >>> # Create a batching mode
+        >>> batching = brainstate.mixin.Batching(batch_size=32, batch_axis=0)
+        >>> print(batching)  # Batching(in_size=32, axis=0)
+        >>>
+        >>> # Access batch parameters
+        >>> print(f"Processing {batching.batch_size} samples at once")
+        >>> print(f"Batch dimension is axis {batching.batch_axis}")
+
+    Using in a model:
+
+    .. code-block:: python
+
+        >>> import jax.numpy as jnp
+        >>>
+        >>> class BatchedModel:
+        ...     def __init__(self):
+        ...         self.mode = None
+        ...
+        ...     def set_batch_mode(self, batch_size, batch_axis=0):
+        ...         self.mode = brainstate.mixin.Batching(batch_size, batch_axis)
+        ...
+        ...     def process(self, x):
+        ...         if self.mode is not None and self.mode.has(brainstate.mixin.Batching):
+        ...             # Process in batches
+        ...             batch_size = self.mode.batch_size
+        ...             axis = self.mode.batch_axis
+        ...             return jnp.mean(x, axis=axis, keepdims=True)
+        ...         return x
+        >>>
+        >>> model = BatchedModel()
+        >>> model.set_batch_mode(batch_size=64)
+        >>>
+        >>> # Process batched data
+        >>> data = jnp.random.randn(64, 100)  # 64 samples, 100 features
+        >>> result = model.process(data)
+
+    Combining with other modes:
+
+    .. code-block:: python
+
+        >>> # Combine batching with training mode
+        >>> training = brainstate.mixin.Training()
+        >>> batching = brainstate.mixin.Batching(batch_size=128)
+        >>> combined = brainstate.mixin.JointMode(training, batching)
+        >>>
+        >>> # Use in a training loop
+        >>> def train_step(model, data, mode):
+        ...     if mode.has(brainstate.mixin.Batching):
+        ...         # Split data into batches
+        ...         batch_size = mode.batch_size
+        ...         # ... batched processing ...
+        ...     if mode.has(brainstate.mixin.Training):
+        ...         # Apply training-specific operations
+        ...         # ... training logic ...
+        ...     pass
+    """
 
     def __init__(self, batch_size: int = 1, batch_axis: int = 0):
         self.batch_size = batch_size
         self.batch_axis = batch_axis
 
     def __repr__(self):
+        """
+        String representation showing batch configuration.
+
+        Returns
+        -------
+        str
+            A string showing the batch size and axis.
+        """
         return f'{self.__class__.__name__}(in_size={self.batch_size}, axis={self.batch_axis})'
 
 
 class Training(Mode):
-    """Training mode."""
+    """
+    Mode indicating training computation.
+
+    This mode specifies that the model is in training mode, which typically
+    enables behaviors like dropout, batch normalization in training mode,
+    gradient computation, etc.
+
+    Examples
+    --------
+    Basic training mode:
+
+    .. code-block:: python
+
+        >>> import brainstate
+        >>>
+        >>> # Create training mode
+        >>> training = brainstate.mixin.Training()
+        >>> print(training)  # Training
+        >>>
+        >>> # Check mode
+        >>> assert training.is_a(brainstate.mixin.Training)
+        >>> assert training.has(brainstate.mixin.Mode)
+
+    Using in a model with dropout:
+
+    .. code-block:: python
+
+        >>> import brainstate
+        >>> import jax
+        >>> import jax.numpy as jnp
+        >>>
+        >>> class ModelWithDropout:
+        ...     def __init__(self, dropout_rate=0.5):
+        ...         self.dropout_rate = dropout_rate
+        ...         self.mode = None
+        ...
+        ...     def set_training(self, is_training=True):
+        ...         if is_training:
+        ...             self.mode = brainstate.mixin.Training()
+        ...         else:
+        ...             self.mode = brainstate.mixin.Mode()  # Evaluation mode
+        ...
+        ...     def forward(self, x, rng_key):
+        ...         # Apply dropout only during training
+        ...         if self.mode is not None and self.mode.has(brainstate.mixin.Training):
+        ...             keep_prob = 1.0 - self.dropout_rate
+        ...             mask = jax.random.bernoulli(rng_key, keep_prob, x.shape)
+        ...             x = jnp.where(mask, x / keep_prob, 0)
+        ...         return x
+        >>>
+        >>> model = ModelWithDropout()
+        >>>
+        >>> # Training mode
+        >>> model.set_training(True)
+        >>> key = jax.random.PRNGKey(0)
+        >>> x_train = jnp.ones((10, 20))
+        >>> out_train = model.forward(x_train, key)  # Dropout applied
+        >>>
+        >>> # Evaluation mode
+        >>> model.set_training(False)
+        >>> out_eval = model.forward(x_train, key)  # No dropout
+
+    Combining with batching:
+
+    .. code-block:: python
+
+        >>> # Create combined training and batching mode
+        >>> training = brainstate.mixin.Training()
+        >>> batching = brainstate.mixin.Batching(batch_size=32)
+        >>> mode = brainstate.mixin.JointMode(training, batching)
+        >>>
+        >>> # Use in training configuration
+        >>> class Trainer:
+        ...     def __init__(self, model, mode):
+        ...         self.model = model
+        ...         self.mode = mode
+        ...
+        ...     def train_epoch(self, data):
+        ...         if self.mode.has(brainstate.mixin.Training):
+        ...             # Enable training-specific behaviors
+        ...             self.model.set_training(True)
+        ...
+        ...         if self.mode.has(brainstate.mixin.Batching):
+        ...             # Process in batches
+        ...             batch_size = self.mode.batch_size
+        ...             # ... batched training loop ...
+        ...         pass
+    """
     pass