brainstate 0.1.0.post20250503__py2.py3-none-any.whl → 0.1.2__py2.py3-none-any.whl

brainstate/random/_rand_seed.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from __future__ import annotations
 
 from contextlib import contextmanager
 from typing import Optional

brainstate/random/_rand_seed_test.py

@@ -13,14 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import unittest
 
 import jax.numpy as jnp
 import jax.random
 
-import brainstate as bst
+import brainstate
 
 
 class TestRandom(unittest.TestCase):
@@ -28,23 +26,23 @@ class TestRandom(unittest.TestCase):
     def test_seed2(self):
         test_seed = 299
         key = jax.random.PRNGKey(test_seed)
-        bst.random.seed(key)
+        brainstate.random.seed(key)
 
         @jax.jit
         def jit_seed(key):
-            bst.random.seed(key)
-            with bst.random.seed_context(key):
-                print(bst.random.DEFAULT.value)
+            brainstate.random.seed(key)
+            with brainstate.random.seed_context(key):
+                print(brainstate.random.DEFAULT.value)
 
         jit_seed(key)
         jit_seed(1)
         jit_seed(None)
-        bst.random.seed(1)
+        brainstate.random.seed(1)
 
     def test_seed(self):
         test_seed = 299
-        bst.random.seed(test_seed)
-        a = bst.random.rand(3)
-        bst.random.seed(test_seed)
-        b = bst.random.rand(3)
+        brainstate.random.seed(test_seed)
+        a = brainstate.random.rand(3)
+        brainstate.random.seed(test_seed)
+        b = brainstate.random.rand(3)
         self.assertTrue(jnp.array_equal(a, b))

brainstate/random/_rand_state.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 
 # -*- coding: utf-8 -*-
-from __future__ import annotations
 
 from functools import partial
 from operator import index

brainstate/surrogate.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 
 # -*- coding: utf-8 -*-
-from __future__ import annotations
 
 import jax
 import jax.numpy as jnp

brainstate/typing.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import builtins
 import functools as ft
 import importlib

brainstate/util/_caller.py

@@ -15,8 +15,6 @@
 # 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
 
@@ -82,18 +80,18 @@ class CallableProxy:
     def __call__(self, *args, **kwargs):
         return self._callable(self._accessor, *args, **kwargs)
 
-    def __getattr__(self, name) -> CallableProxy:
+    def __getattr__(self, name) -> 'CallableProxy':
         return CallableProxy(self._callable, getattr(self._accessor, name))
 
-    def __getitem__(self, key) -> CallableProxy:
+    def __getitem__(self, key) -> 'CallableProxy':
         return CallableProxy(self._callable, self._accessor[key])
 
 
 class ApplyCaller(Protocol, Generic[A]):
-    def __getattr__(self, __name) -> ApplyCaller[A]:
+    def __getattr__(self, __name) -> 'ApplyCaller[A]':
         ...
 
-    def __getitem__(self, __name) -> ApplyCaller[A]:
+    def __getitem__(self, __name) -> 'ApplyCaller[A]':
         ...
 
     def __call__(self, *args, **kwargs) -> tuple[Any, A]:

brainstate/util/_others.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import copy
 import functools
 import gc