update annotation

braintaichi/_eventop/__init__.py

@@ -20,4 +20,3 @@
 __all__ = _main_all
 
 del _main_all
-

braintaichi/_eventop/_event_csrmm.py

@@ -15,6 +15,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 from typing import Tuple
 
 import brainunit as u
@@ -30,6 +32,7 @@
 from braintaichi._sparseop._sparse_csrmm import raw_csrmm_taichi as normal_csrmm
 from braintaichi._sparseop._sparse_utils import csr_to_coo
 
+
 def raw_event_csrmm_taichi(
     data: jax.typing.ArrayLike | u.Quantity,
     indices: jax.typing.ArrayLike,

braintaichi/_eventop/_event_csrmv.py

@@ -25,11 +25,12 @@
 
 """
 
+from __future__ import annotations
+
 from typing import Union, Tuple
 
 import jax
 import jax.numpy as jnp
-import numpy as np
 import taichi as ti
 from jax.interpreters import ad
 

braintaichi/_eventop/main.py

@@ -14,6 +14,8 @@
 # ==============================================================================
 
 
+from __future__ import annotations
+
 from typing import Union, Tuple
 
 import brainunit as u

braintaichi/_jitconnop/__init__.py

@@ -14,10 +14,10 @@
 # ==============================================================================
 
 
-from .main import *
-from .main import __all__ as _main_all
 from ._taichi_rand import *
 from ._taichi_rand import __all__ as _taichi_rand_all
+from .main import *
+from .main import __all__ as _main_all
 
 __all__ = _main_all + _taichi_rand_all
 

braintaichi/_jitconnop/_jit_csrmv.py

@@ -15,6 +15,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 import numbers
 from typing import Tuple, Optional
 
@@ -24,8 +26,8 @@
 from jax import numpy as jnp
 from jax.interpreters import ad
 
-from braintaichi._primitive._xla_custom_op import XLACustomOp
 from braintaichi._misc import _get_dtype, set_module_as
+from braintaichi._primitive._xla_custom_op import XLACustomOp
 from ._taichi_rand import (lfsr88_key, lfsr88_random_integers, lfsr88_uniform, lfsr88_normal)
 
 __all__ = [
@@ -35,7 +37,6 @@
 ]
 
 
-
 @set_module_as('braintaichi')
 def get_homo_weight_matrix(
     weight: float,

braintaichi/_jitconnop/_jit_event_csrmv.py

@@ -16,6 +16,8 @@
 # -*- coding: utf-8 -*-
 
 
+from __future__ import annotations
+
 from typing import Tuple
 
 import jax

braintaichi/_jitconnop/_taichi_rand.py

@@ -15,6 +15,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 import brainstate as bst
 import jax
 import taichi as ti

braintaichi/_jitconnop/main.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
+
 from typing import Tuple, Optional
 
 import jax

braintaichi/_misc.py

@@ -15,6 +15,7 @@
 
 import jax
 
+
 def _get_dtype(v):
   if hasattr(v, 'dtype'):
     dtype = v.dtype
@@ -23,12 +24,9 @@ def _get_dtype(v):
   return dtype
 
 
-
 def set_module_as(name: str):
   def decorator(f):
     f.__module__ = name
     return f
 
   return decorator
-
-

braintaichi/_primitive/_ad_support.py

@@ -15,6 +15,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 import functools
 from functools import partial
 

braintaichi/_primitive/_batch_utils.py

@@ -15,6 +15,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 from functools import partial
 
 import jax.numpy as jnp

braintaichi/_primitive/_mlir_translation_rule.py

@@ -15,6 +15,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 import contextlib
 import hashlib
 import inspect
@@ -37,7 +39,7 @@
 
 # --- REGISTER CUSTOM CALL TARGETS on CPU platforms ###
 try:
-  from braintaichi import cpu_ops # noqa
+  from braintaichi import cpu_ops  # noqa
 
   for _name, _value in cpu_ops.registrations().items():
     xla_client.register_custom_call_target(_name, _value, platform="cpu")
@@ -46,7 +48,7 @@
 
 # --- REGISTER CUSTOM CALL TARGETS on GPU platforms ###
 try:
-  from braintaichi import gpu_ops # noqa
+  from braintaichi import gpu_ops  # noqa
 
   for _name, _value in gpu_ops.registrations().items():
     xla_client.register_custom_call_target(_name, _value, platform="gpu")

braintaichi/_primitive/_xla_custom_op.py

@@ -15,6 +15,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 from functools import partial
 from typing import Callable, Sequence, Tuple, Protocol, Optional, Union
 

braintaichi/_sparseop/_sparse_coomv.py

@@ -15,6 +15,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 import warnings
 from functools import partial
 

braintaichi/_sparseop/_sparse_csrmm.py

@@ -15,6 +15,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 from typing import Tuple
 
 import brainunit as u

braintaichi/_sparseop/_sparse_csrmv.py

@@ -15,6 +15,9 @@
 
 # -*- coding: utf-8 -*-
 
+
+from __future__ import annotations
+
 from typing import Tuple
 
 import brainunit as u

braintaichi/_sparseop/_sparse_utils.py

@@ -15,10 +15,11 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 import warnings
 from typing import Tuple
 
-import jax
 import numpy as np
 from jax import core, numpy as jnp
 from jax.interpreters import mlir, ad
@@ -33,7 +34,6 @@
 ]
 
 
-
 def coo_to_csr(
     pre_ids: jnp.ndarray,
     post_ids: jnp.ndarray,

braintaichi/_sparseop/main.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 # ==============================================================================
 
+
+from __future__ import annotations
+
 from typing import Tuple
 
 import brainunit as u