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

brainstate/__init__.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,42 +17,153 @@
 A ``State``-based Transformation System for Program Compilation and Augmentation
 """
 
-__version__ = "0.1.9"
+__version__ = "0.2.0"
+__versio_info__ = (0, 2, 0)
 
-from . import augment
-from . import compile
 from . import environ
-from . import functional
 from . import graph
-from . import init
 from . import mixin
 from . import nn
-from . import optim
 from . import random
-from . import surrogate
 from . import transform
 from . import typing
 from . import util
+from ._error import *
+from ._error import __all__ as _error_all
 from ._state import *
 from ._state import __all__ as _state_all
 
+# Create deprecated module proxies with scoped APIs
+from ._deprecation import create_deprecated_module_proxy
+
+# Augment module scope
+_augment_apis = {
+    'GradientTransform': 'brainstate.transform._autograd',
+    'grad': 'brainstate.transform._autograd',
+    'vector_grad': 'brainstate.transform._autograd',
+    'hessian': 'brainstate.transform._autograd',
+    'jacobian': 'brainstate.transform._autograd',
+    'jacrev': 'brainstate.transform._autograd',
+    'jacfwd': 'brainstate.transform._autograd',
+    'abstract_init': 'brainstate.transform._eval_shape',
+    'vmap': 'brainstate.transform._mapping',
+    'pmap': 'brainstate.transform._mapping',
+    'map': 'brainstate.transform._mapping',
+    'vmap_new_states': 'brainstate.transform._mapping',
+    'restore_rngs': 'brainstate.transform._random',
+}
+
+augment = create_deprecated_module_proxy(
+    deprecated_name='brainstate.augment',
+    replacement_module=transform,
+    replacement_name='brainstate.transform',
+    scoped_apis=_augment_apis
+)
+
+# Compile module scope
+_compile_apis = {
+    'checkpoint': 'brainstate.transform._ad_checkpoint',
+    'remat': 'brainstate.transform._ad_checkpoint',
+    'cond': 'brainstate.transform._conditions',
+    'switch': 'brainstate.transform._conditions',
+    'ifelse': 'brainstate.transform._conditions',
+    'jit_error_if': 'brainstate.transform._error_if',
+    'jit': 'brainstate.transform._jit',
+    'scan': 'brainstate.transform._loop_collect_return',
+    'checkpointed_scan': 'brainstate.transform._loop_collect_return',
+    'for_loop': 'brainstate.transform._loop_collect_return',
+    'checkpointed_for_loop': 'brainstate.transform._loop_collect_return',
+    'while_loop': 'brainstate.transform._loop_no_collection',
+    'bounded_while_loop': 'brainstate.transform._loop_no_collection',
+    'StatefulFunction': 'brainstate.transform._make_jaxpr',
+    'make_jaxpr': 'brainstate.transform._make_jaxpr',
+    'ProgressBar': 'brainstate.transform._progress_bar',
+}
+
+compile = create_deprecated_module_proxy(
+    deprecated_name='brainstate.compile',
+    replacement_module=transform,
+    replacement_name='brainstate.transform',
+    scoped_apis=_compile_apis
+)
+
+# Functional module scope - use direct attribute access from nn module
+_functional_apis = {
+    'weight_standardization': 'brainstate.nn._normalizations',
+    'clip_grad_norm': 'brainstate.nn._others',
+    'tanh': 'brainstate.nn._activations',
+    'relu': 'brainstate.nn._activations',
+    'squareplus': 'brainstate.nn._activations',
+    'softplus': 'brainstate.nn._activations',
+    'soft_sign': 'brainstate.nn._activations',
+    'sigmoid': 'brainstate.nn._activations',
+    'silu': 'brainstate.nn._activations',
+    'swish': 'brainstate.nn._activations',
+    'log_sigmoid': 'brainstate.nn._activations',
+    'elu': 'brainstate.nn._activations',
+    'leaky_relu': 'brainstate.nn._activations',
+    'hard_tanh': 'brainstate.nn._activations',
+    'celu': 'brainstate.nn._activations',
+    'selu': 'brainstate.nn._activations',
+    'gelu': 'brainstate.nn._activations',
+    'glu': 'brainstate.nn._activations',
+    'logsumexp': 'brainstate.nn._activations',
+    'log_softmax': 'brainstate.nn._activations',
+    'softmax': 'brainstate.nn._activations',
+    'standardize': 'brainstate.nn._activations',
+    'relu6': 'brainstate.nn._activations',
+    'hard_sigmoid': 'brainstate.nn._activations',
+    'sparse_plus': 'brainstate.nn._activations',
+    'hard_silu': 'brainstate.nn._activations',
+    'hard_swish': 'brainstate.nn._activations',
+    'hard_shrink': 'brainstate.nn._activations',
+    'rrelu': 'brainstate.nn._activations',
+    'mish': 'brainstate.nn._activations',
+    'soft_shrink': 'brainstate.nn._activations',
+    'prelu': 'brainstate.nn._activations',
+    'softmin': 'brainstate.nn._activations',
+    'one_hot': 'brainstate.nn._activations',
+    'sparse_sigmoid': 'brainstate.nn._activations',
+}
+
+functional = create_deprecated_module_proxy(
+    deprecated_name='brainstate.functional',
+    replacement_module=nn,
+    replacement_name='brainstate.nn',
+    scoped_apis=_functional_apis
+)
+
+
+def __getattr__(name):
+    if name in ['surrogate', 'init', 'optim']:
+        import warnings
+        warnings.warn(
+            f"brainstate.{name} module is deprecated and will be removed in a future version. "
+            f"Please use braintools.{name} instead.",
+            DeprecationWarning,
+            stacklevel=2
+        )
+        import braintools
+        return getattr(braintools, name)
+    raise AttributeError(
+        f'module {__name__!r} has no attribute {name!r}'
+    )
+
+
 __all__ = [
-    'augment',
-    'compile',
     'environ',
-    'functional',
     'graph',
-    'init',
     'mixin',
     'nn',
-    'optim',
     'random',
-    'surrogate',
+    'transform',
     'typing',
     'util',
-    'transform',
+    # Deprecated modules
+    'augment',
+    'compile',
+    'functional',
 ]
-__all__ = __all__ + _state_all
-
-# ----------------------- #
-del _state_all
+__all__ = __all__ + _state_all + _error_all
+del _state_all, create_deprecated_module_proxy, _augment_apis, _compile_apis, _functional_apis
+del _error_all

brainstate/_compatible_import.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.
@@ -15,12 +15,40 @@
 
 # -*- coding: utf-8 -*-
 
+"""
+Compatibility layer for JAX version differences.
+
+This module provides a compatibility layer to handle differences between various
+versions of JAX, ensuring that BrainState works correctly across different JAX
+versions. It imports the appropriate modules and functions based on the detected
+JAX version and provides fallback implementations when necessary.
+
+Key Features:
+    - Version-aware imports for JAX core functionality
+    - Compatibility wrappers for changed APIs
+    - Fallback implementations for deprecated functions
+    - Type-safe utility functions
+
+Examples:
+    Basic usage:
+
+    >>> from brainstate._compatible_import import safe_map, safe_zip
+    >>> result = safe_map(lambda x: x * 2, [1, 2, 3])
+    >>> pairs = safe_zip([1, 2, 3], ['a', 'b', 'c'])
+
+    Using JAX core types:
+
+    >>> from brainstate._compatible_import import Primitive, ClosedJaxpr
+    >>> # These imports work across different JAX versions
+"""
 
 from contextlib import contextmanager
 from functools import partial
 from typing import Iterable, Hashable, TypeVar, Callable
 
 import jax
+from jax.core import get_aval, Tracer
+from saiunit._compatible_import import wrap_init
 
 __all__ = [
     'ClosedJaxpr',
@@ -36,6 +64,12 @@ __all__ = [
     'wraps',
     'Device',
     'wrap_init',
+    'Var',
+    'JaxprEqn',
+    'Jaxpr',
+    'Literal',
+
+    'make_iota', 'to_elt', 'BatchTracer', 'BatchTrace',
 ]
 
 T = TypeVar("T")
@@ -43,24 +77,45 @@ T1 = TypeVar("T1")
 T2 = TypeVar("T2")
 T3 = TypeVar("T3")
 
-from saiunit._compatible_import import wrap_init
-
-from jax.core import get_aval, Tracer
-
 if jax.__version_info__ < (0, 5, 0):
     from jax.lib.xla_client import Device
 else:
     from jax import Device
 
+if jax.__version_info__ < (0, 7, 1):
+    from jax.interpreters.batching import make_iota, to_elt, BatchTracer, BatchTrace
+else:
+    from jax._src.interpreters.batching import make_iota, to_elt, BatchTracer, BatchTrace
+
 if jax.__version_info__ < (0, 4, 38):
     from jax.core import ClosedJaxpr, extend_axis_env_nd, Primitive, jaxpr_as_fun
+    from jax.core import Primitive, Var, JaxprEqn, Jaxpr, ClosedJaxpr, Literal
 else:
     from jax.extend.core import ClosedJaxpr, Primitive, jaxpr_as_fun
+    from jax.extend.core import Primitive, Var, JaxprEqn, Jaxpr, ClosedJaxpr, Literal
     from jax.core import trace_ctx
 
 
     @contextmanager
     def extend_axis_env_nd(name_size_pairs: Iterable[tuple[Hashable, int]]):
+        """
+        Context manager to temporarily extend the JAX axis environment.
+
+        Extends the current JAX axis environment with new named axes for
+        vectorized computations, then restores the previous environment.
+
+        Args:
+            name_size_pairs: Iterable of (name, size) tuples specifying
+                           the named axes to add to the environment.
+
+        Yields:
+            None: Context with extended axis environment.
+
+        Examples:
+            >>> with extend_axis_env_nd([('batch', 32), ('seq', 128)]):
+            ...     # Code using vectorized operations with named axes
+            ...     pass
+        """
         prev = trace_ctx.axis_env
         try:
             trace_ctx.set_axis_env(prev.extend_pure(name_size_pairs))
@@ -73,6 +128,29 @@ if jax.__version_info__ < (0, 6, 0):
 
 else:
     def safe_map(f, *args):
+        """
+        Map a function over multiple sequences with length checking.
+
+        Applies a function to corresponding elements from multiple sequences,
+        ensuring all sequences have the same length.
+
+        Args:
+            f: Function to apply to elements from each sequence.
+            *args: Variable number of sequences to map over.
+
+        Returns:
+            list: Results of applying f to corresponding elements.
+
+        Raises:
+            AssertionError: If input sequences have different lengths.
+
+        Examples:
+            >>> safe_map(lambda x, y: x + y, [1, 2, 3], [4, 5, 6])
+            [5, 7, 9]
+
+            >>> safe_map(str.upper, ['a', 'b', 'c'])
+            ['A', 'B', 'C']
+        """
         args = list(map(list, args))
         n = len(args[0])
         for arg in args[1:]:
@@ -81,6 +159,28 @@ else:
 
 
     def safe_zip(*args):
+        """
+        Zip multiple sequences with length checking.
+
+        Combines corresponding elements from multiple sequences into tuples,
+        ensuring all sequences have the same length.
+
+        Args:
+            *args: Variable number of sequences to zip together.
+
+        Returns:
+            list: List of tuples containing corresponding elements.
+
+        Raises:
+            AssertionError: If input sequences have different lengths.
+
+        Examples:
+            >>> safe_zip([1, 2, 3], ['a', 'b', 'c'])
+            [(1, 'a'), (2, 'b'), (3, 'c')]
+
+            >>> safe_zip([1, 2], [3, 4], [5, 6])
+            [(1, 3, 5), (2, 4, 6)]
+        """
         args = list(map(list, args))
         n = len(args[0])
         for arg in args[1:]:
@@ -89,7 +189,32 @@ else:
 
 
     def unzip2(xys: Iterable[tuple[T1, T2]]) -> tuple[tuple[T1, ...], tuple[T2, ...]]:
-        """Unzip sequence of length-2 tuples into two tuples."""
+        """
+        Unzip sequence of length-2 tuples into two tuples.
+
+        Takes an iterable of 2-tuples and separates them into two tuples
+        containing the first and second elements respectively.
+
+        Args:
+            xys: Iterable of 2-tuples to unzip.
+
+        Returns:
+            tuple: A 2-tuple containing:
+                - Tuple of all first elements
+                - Tuple of all second elements
+
+        Examples:
+            >>> pairs = [(1, 'a'), (2, 'b'), (3, 'c')]
+            >>> nums, letters = unzip2(pairs)
+            >>> nums
+            (1, 2, 3)
+            >>> letters
+            ('a', 'b', 'c')
+
+        Notes:
+            We deliberately don't use zip(*xys) because it is lazily evaluated,
+            is too permissive about inputs, and does not guarantee a length-2 output.
+        """
         # Note: we deliberately don't use zip(*xys) because it is lazily evaluated,
         # is too permissive about inputs, and does not guarantee a length-2 output.
         xs: list[T1] = []
@@ -101,6 +226,30 @@ else:
 
 
     def fun_name(fun: Callable):
+        """
+        Extract the name of a function, handling special cases.
+
+        Attempts to get the name of a function, with special handling for
+        partial functions and fallback for unnamed functions.
+
+        Args:
+            fun: The function to get the name from.
+
+        Returns:
+            str: The function name, or "<unnamed function>" if no name available.
+
+        Examples:
+            >>> def my_function():
+            ...     pass
+            >>> fun_name(my_function)
+            'my_function'
+
+            >>> from functools import partial
+            >>> add = lambda x, y: x + y
+            >>> add_one = partial(add, 1)
+            >>> fun_name(add_one)
+            '<lambda>'
+        """
         name = getattr(fun, "__name__", None)
         if name is not None:
             return name
@@ -117,8 +266,34 @@ else:
         **kwargs,
     ) -> Callable[[T], T]:
         """
-        Like functools.wraps, but with finer-grained control over the name and docstring
-        of the resulting function.
+        Enhanced function wrapper with fine-grained control.
+
+        Like functools.wraps, but provides more control over the name and docstring
+        of the resulting function. Useful for creating custom decorators.
+
+        Args:
+            wrapped: The function being wrapped.
+            namestr: Optional format string for the wrapper function name.
+                    Can use {fun} placeholder for the original function name.
+            docstr: Optional format string for the wrapper function docstring.
+                   Can use {fun}, {doc}, and other kwargs as placeholders.
+            **kwargs: Additional keyword arguments for format string substitution.
+
+        Returns:
+            Callable: A decorator function that applies the wrapping.
+
+        Examples:
+            >>> def my_decorator(func):
+            ...     @wraps(func, namestr="decorated_{fun}")
+            ...     def wrapper(*args, **kwargs):
+            ...         return func(*args, **kwargs)
+            ...     return wrapper
+
+            >>> @my_decorator
+            ... def example():
+            ...     pass
+            >>> example.__name__
+            'decorated_example'
         """
 
         def wrapper(fun: T) -> T:
@@ -141,8 +316,25 @@ else:
 
 
 def to_concrete_aval(aval):
+    """
+    Convert an abstract value to its concrete representation.
+
+    Takes an abstract value and attempts to convert it to a concrete value,
+    handling JAX Tracer objects appropriately.
+
+    Args:
+        aval: The abstract value to convert.
+
+    Returns:
+        The concrete value representation, or the original aval if already concrete.
+
+    Examples:
+        >>> import jax.numpy as jnp
+        >>> arr = jnp.array([1, 2, 3])
+        >>> concrete = to_concrete_aval(arr)
+        # Returns the concrete array value
+    """
     aval = get_aval(aval)
     if isinstance(aval, Tracer):
         return aval.to_concrete_value()
     return aval
-