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

brainstate/util/__init__.py

@@ -1,50 +1,50 @@
-# 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 . import filter
-from .error import *
-from .error import __all__ as _error_all
-from .others import *
-from .others import __all__ as _others_all
-from .pretty_pytree import *
-from .pretty_pytree import __all__ as _mapping_all
-from .pretty_repr import *
-from .pretty_repr import __all__ as _pretty_repr_all
-from .pretty_table import *
-from .pretty_table import __all__ as _table_all
-from .scaling import *
-from .scaling import __all__ as _mem_scale_all
-from .struct import *
-from .struct import __all__ as _struct_all
-
-__all__ = (
-    ['filter']
-    + _others_all
-    + _mem_scale_all
-    + _pretty_repr_all
-    + _struct_all
-    + _error_all
-    + _mapping_all
-    + _table_all
-)
-del (
-    _others_all,
-    _mem_scale_all,
-    _pretty_repr_all,
-    _struct_all,
-    _error_all,
-    _mapping_all,
-    _table_all,
-)
+# 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 . import filter
+from .error import *
+from .error import __all__ as _error_all
+from .others import *
+from .others import __all__ as _others_all
+from .pretty_pytree import *
+from .pretty_pytree import __all__ as _mapping_all
+from .pretty_repr import *
+from .pretty_repr import __all__ as _pretty_repr_all
+from .pretty_table import *
+from .pretty_table import __all__ as _table_all
+from .scaling import *
+from .scaling import __all__ as _mem_scale_all
+from .struct import *
+from .struct import __all__ as _struct_all
+
+__all__ = (
+    ['filter']
+    + _others_all
+    + _mem_scale_all
+    + _pretty_repr_all
+    + _struct_all
+    + _error_all
+    + _mapping_all
+    + _table_all
+)
+del (
+    _others_all,
+    _mem_scale_all,
+    _pretty_repr_all,
+    _struct_all,
+    _error_all,
+    _mapping_all,
+    _table_all,
+)

brainstate/util/caller.py

@@ -1,98 +1,98 @@
-# 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.
-#
-# 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.
-
-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]:
-        ...
+# 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.
+#
+# 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.
+
+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]:
+        ...

brainstate/util/error.py

@@ -1,55 +1,55 @@
-# 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.
-# ==============================================================================
-
-
-__all__ = [
-    'BrainStateError',
-    'TraceContextError',
-]
-
-
-class BrainStateError(Exception):
-    """
-    A custom exception class for BrainState-related errors.
-
-    This exception is raised when a BrainState-specific error occurs during
-    the execution of the program. It serves as a base class for more specific
-    BrainState exceptions.
-
-    Attributes:
-        Inherits all attributes from the built-in Exception class.
-
-    Usage::
-
-        raise BrainStateError("A BrainState-specific error occurred.")
-    """
-    pass
-
-
-class TraceContextError(BrainStateError):
-    """
-    A custom exception class for trace context-related errors in BrainState.
-
-    This exception is raised when an error occurs specifically related to
-    trace context operations or manipulations within the BrainState framework.
-
-    Attributes:
-        Inherits all attributes from the BrainStateError class.
-
-    Usage::
-
-        raise TraceContextError("An error occurred while handling trace context.")
-    """
-    pass
+# 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.
+# ==============================================================================
+
+
+__all__ = [
+    'BrainStateError',
+    'TraceContextError',
+]
+
+
+class BrainStateError(Exception):
+    """
+    A custom exception class for BrainState-related errors.
+
+    This exception is raised when a BrainState-specific error occurs during
+    the execution of the program. It serves as a base class for more specific
+    BrainState exceptions.
+
+    Attributes:
+        Inherits all attributes from the built-in Exception class.
+
+    Usage::
+
+        raise BrainStateError("A BrainState-specific error occurred.")
+    """
+    pass
+
+
+class TraceContextError(BrainStateError):
+    """
+    A custom exception class for trace context-related errors in BrainState.
+
+    This exception is raised when an error occurs specifically related to
+    trace context operations or manipulations within the BrainState framework.
+
+    Attributes:
+        Inherits all attributes from the BrainStateError class.
+
+    Usage::
+
+        raise TraceContextError("An error occurred while handling trace context.")
+    """
+    pass