brainstate 0.0.2.post20241010__py2.py3-none-any.whl → 0.1.0.post20241122__py2.py3-none-any.whl

brainstate/{nn/_projection/_utils.py → transform.py}

@@ -13,5 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-def is_instance(x, cls):
-  assert isinstance(x, cls), 'The input should be an instance of {}!'.format(cls)
+# This module is going to be deleted in the future (near 2025-06).

brainstate/typing.py

@@ -13,73 +13,96 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
 
+import builtins
 import functools as ft
+import importlib
 import inspect
-import typing
-from typing import Sequence, Protocol, Union, Any, Generic, TypeVar, Tuple
 
-import brainunit as bu
+import brainunit as u
 import jax
 import numpy as np
 
+tp = importlib.import_module("typing")
+
 __all__ = [
-  'PyTree',
-  'Size',
-  'Axes',
-  'SeedOrKey',
-  'ArrayLike',
-  'DType',
-  'DTypeLike',
+    'PathParts',
+    'Predicate',
+    'Filter',
+    'PyTree',
+    'Size',
+    'Axes',
+    'SeedOrKey',
+    'ArrayLike',
+    'DType',
+    'DTypeLike',
+    'Missing',
 ]
 
-_T = TypeVar("_T")
+K = tp.TypeVar('K')
+
+
+@tp.runtime_checkable
+class Key(tp.Hashable, tp.Protocol):
+    def __lt__(self: K, value: K, /) -> bool:
+        ...
+
 
-_Annotation = TypeVar("_Annotation")
+Ellipsis = builtins.ellipsis if tp.TYPE_CHECKING else tp.Any
 
+PathParts = tp.Tuple[Key, ...]
+Predicate = tp.Callable[[PathParts, tp.Any], bool]
+FilterLiteral = tp.Union[type, str, Predicate, bool, Ellipsis, None]
+Filter = tp.Union[FilterLiteral, tp.Tuple['Filter', ...], tp.List['Filter']]
 
-class _Array(Generic[_Annotation]):
-  pass
+_T = tp.TypeVar("_T")
+
+_Annotation = tp.TypeVar("_Annotation")
+
+
+class _Array(tp.Generic[_Annotation]):
+    pass
 
 
 _Array.__module__ = "builtins"
 
 
-def _item_to_str(item: Union[str, type, slice]) -> str:
-  if isinstance(item, slice):
-    if item.step is not None:
-      raise NotImplementedError
-    return _item_to_str(item.start) + ": " + _item_to_str(item.stop)
-  elif item is ...:
-    return "..."
-  elif inspect.isclass(item):
-    return item.__name__
-  else:
-    return repr(item)
+def _item_to_str(item: tp.Union[str, type, slice]) -> str:
+    if isinstance(item, slice):
+        if item.step is not None:
+            raise NotImplementedError
+        return _item_to_str(item.start) + ": " + _item_to_str(item.stop)
+    elif item is ...:
+        return "..."
+    elif inspect.isclass(item):
+        return item.__name__
+    else:
+        return repr(item)
 
 
 def _maybe_tuple_to_str(
-    item: Union[str, type, slice, Tuple[Union[str, type, slice], ...]]
+    item: tp.Union[str, type, slice, tp.Tuple[tp.Union[str, type, slice], ...]]
 ) -> str:
-  if isinstance(item, tuple):
-    if len(item) == 0:
-      # Explicit brackets
-      return "()"
+    if isinstance(item, tuple):
+        if len(item) == 0:
+            # Explicit brackets
+            return "()"
+        else:
+            # No brackets
+            return ", ".join([_item_to_str(i) for i in item])
     else:
-      # No brackets
-      return ", ".join([_item_to_str(i) for i in item])
-  else:
-    return _item_to_str(item)
+        return _item_to_str(item)
 
 
 class Array:
-  def __class_getitem__(cls, item):
-    class X:
-      pass
+    def __class_getitem__(cls, item):
+        class X:
+            pass
 
-    X.__module__ = "builtins"
-    X.__qualname__ = _maybe_tuple_to_str(item)
-    return _Array[X]
+        X.__module__ = "builtins"
+        X.__qualname__ = _maybe_tuple_to_str(item)
+        return _Array[X]
 
 
 # Same __module__ trick here again. (So that we get the correct display when
@@ -89,8 +112,8 @@ class Array:
 Array.__module__ = "builtins"
 
 
-class _FakePyTree(Generic[_T]):
-  pass
+class _FakePyTree(tp.Generic[_T]):
+    pass
 
 
 _FakePyTree.__name__ = "PyTree"
@@ -99,84 +122,84 @@ _FakePyTree.__module__ = "builtins"
 
 
 class _MetaPyTree(type):
-  def __call__(self, *args, **kwargs):
-    raise RuntimeError("PyTree cannot be instantiated")
-
-  # Can't return a generic (e.g. _FakePyTree[item]) because generic aliases don't do
-  # the custom __instancecheck__ that we want.
-  # We can't add that __instancecheck__  via subclassing, e.g.
-  # type("PyTree", (Generic[_T],), {}), because dynamic subclassing of typeforms
-  # isn't allowed.
-  # Likewise we can't do types.new_class("PyTree", (Generic[_T],), {}) because that
-  # has __module__ "types", e.g. we get types.PyTree[int].
-  @ft.lru_cache(maxsize=None)
-  def __getitem__(cls, item):
-    if isinstance(item, tuple):
-      if len(item) == 2:
-
-        class X(PyTree):
-          leaftype = item[0]
-          structure = item[1].strip()
-
-        if not isinstance(X.structure, str):
-          raise ValueError(
-            "The structure annotation `struct` in "
-            "`brainstate.typing.PyTree[leaftype, struct]` must be be a string, "
-            f"e.g. `brainstate.typing.PyTree[leaftype, 'T']`. Got '{X.structure}'."
-          )
-        pieces = X.structure.split()
-        if len(pieces) == 0:
-          raise ValueError(
-            "The string `struct` in `brainstate.typing.PyTree[leaftype, struct]` "
-            "cannot be the empty string."
-          )
-        for piece_index, piece in enumerate(pieces):
-          if (piece_index == 0) or (piece_index == len(pieces) - 1):
-            if piece == "...":
-              continue
-          if not piece.isidentifier():
-            raise ValueError(
-              "The string `struct` in "
-              "`brainstate.typing.PyTree[leaftype, struct]` must be be a "
-              "whitespace-separated sequence of identifiers, e.g. "
-              "`brainstate.typing.PyTree[leaftype, 'T']` or "
-              "`brainstate.typing.PyTree[leaftype, 'foo bar']`.\n"
-              "(Here, 'identifier' is used in the same sense as in "
-              "regular Python, i.e. a valid variable name.)\n"
-              f"Got piece '{piece}' in overall structure '{X.structure}'."
-            )
-        name = str(_FakePyTree[item[0]])[:-1] + ', "' + item[1].strip() + '"]'
-      else:
-        raise ValueError(
-          "The subscript `foo` in `brainstate.typing.PyTree[foo]` must either be a "
-          "leaf type, e.g. `PyTree[int]`, or a 2-tuple of leaf and "
-          "structure, e.g. `PyTree[int, 'T']`. Received a tuple of length "
-          f"{len(item)}."
-        )
-    else:
-      name = str(_FakePyTree[item])
-
-      class X(PyTree):
-        leaftype = item
-        structure = None
-
-    X.__name__ = name
-    X.__qualname__ = name
-    if getattr(typing, "GENERATING_DOCUMENTATION", False):
-      X.__module__ = "builtins"
-    else:
-      X.__module__ = "brainstate.typing"
-    return X
+    def __call__(self, *args, **kwargs):
+        raise RuntimeError("PyTree cannot be instantiated")
+
+    # Can't return a generic (e.g. _FakePyTree[item]) because generic aliases don't do
+    # the custom __instancecheck__ that we want.
+    # We can't add that __instancecheck__  via subclassing, e.g.
+    # type("PyTree", (Generic[_T],), {}), because dynamic subclassing of typeforms
+    # isn't allowed.
+    # Likewise we can't do types.new_class("PyTree", (Generic[_T],), {}) because that
+    # has __module__ "types", e.g. we get types.PyTree[int].
+    @ft.lru_cache(maxsize=None)
+    def __getitem__(cls, item):
+        if isinstance(item, tuple):
+            if len(item) == 2:
+
+                class X(PyTree):
+                    leaftype = item[0]
+                    structure = item[1].strip()
+
+                if not isinstance(X.structure, str):
+                    raise ValueError(
+                        "The structure annotation `struct` in "
+                        "`brainstate.typing.PyTree[leaftype, struct]` must be be a string, "
+                        f"e.g. `brainstate.typing.PyTree[leaftype, 'T']`. Got '{X.structure}'."
+                    )
+                pieces = X.structure.split()
+                if len(pieces) == 0:
+                    raise ValueError(
+                        "The string `struct` in `brainstate.typing.PyTree[leaftype, struct]` "
+                        "cannot be the empty string."
+                    )
+                for piece_index, piece in enumerate(pieces):
+                    if (piece_index == 0) or (piece_index == len(pieces) - 1):
+                        if piece == "...":
+                            continue
+                    if not piece.isidentifier():
+                        raise ValueError(
+                            "The string `struct` in "
+                            "`brainstate.typing.PyTree[leaftype, struct]` must be be a "
+                            "whitespace-separated sequence of identifiers, e.g. "
+                            "`brainstate.typing.PyTree[leaftype, 'T']` or "
+                            "`brainstate.typing.PyTree[leaftype, 'foo bar']`.\n"
+                            "(Here, 'identifier' is used in the same sense as in "
+                            "regular Python, i.e. a valid variable name.)\n"
+                            f"Got piece '{piece}' in overall structure '{X.structure}'."
+                        )
+                name = str(_FakePyTree[item[0]])[:-1] + ', "' + item[1].strip() + '"]'
+            else:
+                raise ValueError(
+                    "The subscript `foo` in `brainstate.typing.PyTree[foo]` must either be a "
+                    "leaf type, e.g. `PyTree[int]`, or a 2-tuple of leaf and "
+                    "structure, e.g. `PyTree[int, 'T']`. Received a tuple of length "
+                    f"{len(item)}."
+                )
+        else:
+            name = str(_FakePyTree[item])
+
+            class X(PyTree):
+                leaftype = item
+                structure = None
+
+        X.__name__ = name
+        X.__qualname__ = name
+        if getattr(tp, "GENERATING_DOCUMENTATION", False):
+            X.__module__ = "builtins"
+        else:
+            X.__module__ = "brainstate.typing"
+        return X
 
 
 # Can't do `class PyTree(Generic[_T]): ...` because we need to override the
 # instancecheck for PyTree[foo], but subclassing
 # `type(Generic[int])`, i.e. `typing._GenericAlias` is disallowed.
 PyTree = _MetaPyTree("PyTree", (), {})
-if getattr(typing, "GENERATING_DOCUMENTATION", False):
-  PyTree.__module__ = "builtins"
+if getattr(tp, "GENERATING_DOCUMENTATION", False):
+    PyTree.__module__ = "builtins"
 else:
-  PyTree.__module__ = "brainstate.typing"
+    PyTree.__module__ = "brainstate.typing"
 PyTree.__doc__ = """Represents a PyTree.
 
 Annotations of the following sorts are supported:
@@ -231,9 +254,9 @@ f. A structure can end with a `...`, to denote that the PyTree must be a prefix
     cases, all named pieces must already have been seen and their structures bound.
 """  # noqa: E501
 
-Size = Union[int, Sequence[int]]
-Axes = Union[int, Sequence[int]]
-SeedOrKey = Union[int, jax.Array, np.ndarray]
+Size = tp.Union[int, tp.Sequence[int]]
+Axes = tp.Union[int, tp.Sequence[int]]
+SeedOrKey = tp.Union[int, jax.Array, np.ndarray]
 
 # --- Array --- #
 
@@ -241,12 +264,12 @@ SeedOrKey = Union[int, jax.Array, np.ndarray]
 # standard JAX array (i.e. not including future non-standard array types like
 # KeyArray and BInt). It's different than np.typing.ArrayLike in that it doesn't
 # accept arbitrary sequences, nor does it accept string data.
-ArrayLike = Union[
-  jax.Array,  # JAX array type
-  np.ndarray,  # NumPy array type
-  np.bool_, np.number,  # NumPy scalar types
-  bool, int, float, complex,  # Python scalar types
-  bu.Quantity,  # Quantity
+ArrayLike = tp.Union[
+    jax.Array,  # JAX array type
+    np.ndarray,  # NumPy array type
+    np.bool_, np.number,  # NumPy scalar types
+    bool, int, float, complex,  # Python scalar types
+    u.Quantity,  # Quantity
 ]
 
 # --- Dtype --- #
@@ -255,9 +278,9 @@ ArrayLike = Union[
 DType = np.dtype
 
 
-class SupportsDType(Protocol):
-  @property
-  def dtype(self) -> DType: ...
+class SupportsDType(tp.Protocol):
+    @property
+    def dtype(self) -> DType: ...
 
 
 # DTypeLike is meant to annotate inputs to np.dtype that return
@@ -265,9 +288,13 @@ class SupportsDType(Protocol):
 # because JAX doesn't support objects or structured dtypes.
 # Unlike np.typing.DTypeLike, we exclude None, and instead require
 # explicit annotations when None is acceptable.
-DTypeLike = Union[
-  str,  # like 'float32', 'int32'
-  type[Any],  # like np.float32, np.int32, float, int
-  np.dtype,  # like np.dtype('float32'), np.dtype('int32')
-  SupportsDType,  # like jnp.float32, jnp.int32
+DTypeLike = tp.Union[
+    str,  # like 'float32', 'int32'
+    type[tp.Any],  # like np.float32, np.int32, float, int
+    np.dtype,  # like np.dtype('float32'), np.dtype('int32')
+    SupportsDType,  # like jnp.float32, jnp.int32
 ]
+
+
+class Missing:
+    pass

brainstate/util/__init__.py

@@ -0,0 +1,52 @@
+# 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 ._dict import *
+from ._dict import __all__ as _mapping_all
+from ._error import *
+from ._error import __all__ as _error_all
+from ._filter import *
+from ._filter import __all__ as _filter_all
+from ._others import *
+from ._others import __all__ as _others_all
+from ._pretty_repr import *
+from ._pretty_repr import __all__ as _pretty_repr_all
+from ._scaling import *
+from ._scaling import __all__ as _mem_scale_all
+from ._struct import *
+from ._struct import __all__ as _struct_all
+from ._visualization import *
+from ._visualization import __all__ as _visualization_all
+
+__all__ = (
+    _others_all
+    + _mem_scale_all
+    + _filter_all
+    + _pretty_repr_all
+    + _struct_all
+    + _error_all
+    + _mapping_all
+    + _visualization_all
+)
+del (
+    _others_all,
+    _mem_scale_all,
+    _filter_all,
+    _pretty_repr_all,
+    _struct_all,
+    _error_all,
+    _mapping_all,
+    _visualization_all,
+)

brainstate/util/_caller.py

@@ -0,0 +1,100 @@
+# The file is adapted from the Flax library (https://github.com/google/flax).
+# The credit should go to the Flax authors.
+#
+# Copyright 2024 The Flax Authors & 2024 BDP Ecosystem.
+#
+# 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 __future__ import annotations
+
+import dataclasses
+from typing import Any, TypeVar, Protocol, Generic
+
+import jax
+
+__all__ = [
+    'DelayedAccessor',
+    'CallableProxy',
+    'ApplyCaller',
+]
+
+A = TypeVar('A', covariant=True)  # type: ignore[not-supported-yet]
+
+
+def _identity(x):
+    return x
+
+
+@dataclasses.dataclass(frozen=True)
+class GetItem:
+    key: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class GetAttr:
+    name: str
+
+
+@dataclasses.dataclass(frozen=True)
+class DelayedAccessor:
+    actions: tuple[GetItem | GetAttr, ...] = ()
+
+    def __call__(self, x):
+        for action in self.actions:
+            if isinstance(action, GetItem):
+                x = x[action.key]
+            elif isinstance(action, GetAttr):
+                x = getattr(x, action.name)
+        return x
+
+    def __getattr__(self, name):
+        return DelayedAccessor(self.actions + (GetAttr(name),))
+
+    def __getitem__(self, key):
+        return DelayedAccessor(self.actions + (GetItem(key),))
+
+
+jax.tree_util.register_static(DelayedAccessor)
+
+
+class _AccessorCall(Protocol):
+    def __call__(self, accessor: DelayedAccessor, /, *args, **kwargs) -> Any:
+        ...
+
+
+class CallableProxy:
+    def __init__(
+        self, fun: _AccessorCall, accessor: DelayedAccessor | None = None
+    ):
+        self._callable = fun
+        self._accessor = DelayedAccessor() if accessor is None else accessor
+
+    def __call__(self, *args, **kwargs):
+        return self._callable(self._accessor, *args, **kwargs)
+
+    def __getattr__(self, name) -> CallableProxy:
+        return CallableProxy(self._callable, getattr(self._accessor, name))
+
+    def __getitem__(self, key) -> CallableProxy:
+        return CallableProxy(self._callable, self._accessor[key])
+
+
+class ApplyCaller(Protocol, Generic[A]):
+    def __getattr__(self, __name) -> ApplyCaller[A]:
+        ...
+
+    def __getitem__(self, __name) -> ApplyCaller[A]:
+        ...
+
+    def __call__(self, *args, **kwargs) -> tuple[Any, A]:
+        ...