update codes

chaobrain · Aug 25, 2024 · d49a814 · d49a814
1 parent f7fd900
commit d49a814
Show file tree

Hide file tree

Showing 16 changed files with 12 additions and 45 deletions.
diff --git a/braintaichi/_eventop/_event_csrmm.py b/braintaichi/_eventop/_event_csrmm.py
@@ -15,9 +15,7 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
-from typing import Tuple
+from typing import Tuple, Union
 
 import brainunit as u
 import jax
@@ -34,7 +32,7 @@
 
 
 def raw_event_csrmm_taichi(
-    data: jax.typing.ArrayLike | u.Quantity,
+    data: Union[jax.typing.ArrayLike, u.Quantity],
     indices: jax.typing.ArrayLike,
     indptr: jax.typing.ArrayLike,
     matrix: jax.typing.ArrayLike,

diff --git a/braintaichi/_eventop/_event_csrmv.py b/braintaichi/_eventop/_event_csrmv.py
@@ -25,8 +25,6 @@
 
 """
 
-from __future__ import annotations
-
 from typing import Union, Tuple
 
 import jax

diff --git a/braintaichi/_eventop/main.py b/braintaichi/_eventop/main.py
@@ -14,8 +14,6 @@
 # ==============================================================================
 
 
-from __future__ import annotations
-
 from typing import Union, Tuple
 
 import brainunit as u
@@ -33,7 +31,7 @@
 
 
 def event_csrmm(
-    data: jax.typing.ArrayLike | u.Quantity,
+    data: Union[jax.typing.ArrayLike, u.Quantity],
     indices: jax.typing.ArrayLike,
     indptr: jax.typing.ArrayLike,
     matrix: jax.typing.ArrayLike,
@@ -61,7 +59,7 @@ def event_csrmm(
 
 
 def event_csrmv(
-    data: Union[float, jax.Array],
+    data: Union[jax.typing.ArrayLike, u.Quantity],
     indices: jax.Array,
     indptr: jax.Array,
     events: jax.Array,

diff --git a/braintaichi/_jitconnop/_jit_csrmv.py b/braintaichi/_jitconnop/_jit_csrmv.py
@@ -15,8 +15,6 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 import numbers
 from typing import Tuple, Optional
 

diff --git a/braintaichi/_jitconnop/_jit_event_csrmv.py b/braintaichi/_jitconnop/_jit_event_csrmv.py
@@ -16,8 +16,6 @@
 # -*- coding: utf-8 -*-
 
 
-from __future__ import annotations
-
 from typing import Tuple
 
 import jax

diff --git a/braintaichi/_jitconnop/_taichi_rand.py b/braintaichi/_jitconnop/_taichi_rand.py
@@ -15,8 +15,6 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 import brainstate as bst
 import jax
 import taichi as ti

diff --git a/braintaichi/_jitconnop/main.py b/braintaichi/_jitconnop/main.py
@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 from typing import Tuple, Optional
 
 import jax

diff --git a/braintaichi/_primitive/_ad_support.py b/braintaichi/_primitive/_ad_support.py
@@ -15,8 +15,6 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 import functools
 from functools import partial
 

diff --git a/braintaichi/_primitive/_batch_utils.py b/braintaichi/_primitive/_batch_utils.py
@@ -15,8 +15,6 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 from functools import partial
 
 import jax.numpy as jnp

diff --git a/braintaichi/_primitive/_mlir_translation_rule.py b/braintaichi/_primitive/_mlir_translation_rule.py
@@ -15,8 +15,6 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 import contextlib
 import hashlib
 import inspect

diff --git a/braintaichi/_primitive/_xla_custom_op.py b/braintaichi/_primitive/_xla_custom_op.py
@@ -15,8 +15,6 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 from functools import partial
 from typing import Callable, Sequence, Tuple, Protocol, Optional, Union
 

diff --git a/braintaichi/_sparseop/_sparse_coomv.py b/braintaichi/_sparseop/_sparse_coomv.py
@@ -15,8 +15,6 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 import warnings
 from functools import partial
 

diff --git a/braintaichi/_sparseop/_sparse_csrmm.py b/braintaichi/_sparseop/_sparse_csrmm.py
@@ -15,9 +15,7 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
-from typing import Tuple
+from typing import Tuple, Union
 
 import brainunit as u
 import jax
@@ -32,7 +30,7 @@
 
 
 def raw_csrmm_taichi(
-    data: jax.typing.ArrayLike | u.Quantity,
+    data: Union[jax.typing.ArrayLike, u.Quantity],
     indices: jax.typing.ArrayLike,
     indptr: jax.typing.ArrayLike,
     matrix: jax.typing.ArrayLike,

diff --git a/braintaichi/_sparseop/_sparse_csrmv.py b/braintaichi/_sparseop/_sparse_csrmv.py
@@ -16,9 +16,7 @@
 # -*- coding: utf-8 -*-
 
 
-from __future__ import annotations
-
-from typing import Tuple
+from typing import Tuple, Union
 
 import brainunit as u
 import jax
@@ -33,7 +31,7 @@
 
 
 def raw_csrmv_taichi(
-    data: jax.typing.ArrayLike | u.Quantity,
+    data: Union[jax.typing.ArrayLike, u.Quantity],
     indices: jax.typing.ArrayLike,
     indptr: jax.typing.ArrayLike,
     vector: jax.typing.ArrayLike,

diff --git a/braintaichi/_sparseop/_sparse_utils.py b/braintaichi/_sparseop/_sparse_utils.py
@@ -15,8 +15,6 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 import warnings
 from typing import Tuple
 

diff --git a/braintaichi/_sparseop/main.py b/braintaichi/_sparseop/main.py
@@ -13,10 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-
-from __future__ import annotations
-
-from typing import Tuple
+from typing import Tuple, Union
 
 import brainunit as u
 import jax
@@ -36,7 +33,7 @@
 
 @set_module_as('braintaichi')
 def coomv(
-    data: jax.typing.ArrayLike | u.Quantity,
+    data: Union[jax.typing.ArrayLike, u.Quantity],
     row: jax.typing.ArrayLike,
     col: jax.typing.ArrayLike,
     vector: jax.typing.ArrayLike,
@@ -110,7 +107,7 @@ def coomv(
 
 @set_module_as('braintaichi')
 def csrmm(
-    data: jax.typing.ArrayLike | u.Quantity,
+    data: Union[jax.typing.ArrayLike, u.Quantity],
     indices: jax.typing.ArrayLike,
     indptr: jax.typing.ArrayLike,
     matrix: jax.typing.ArrayLike,
@@ -140,7 +137,7 @@ def csrmm(
 
 @set_module_as('braintaichi')
 def csrmv(
-    data: jax.typing.ArrayLike | u.Quantity,
+    data: Union[jax.typing.ArrayLike, u.Quantity],
     indices: jax.typing.ArrayLike,
     indptr: jax.typing.ArrayLike,
     vector: jax.typing.ArrayLike,
-Original file line number
+Diff line change
@@ Expand Up / @@ -25,8 +25,6 @@ @@
     """
-    from __future__ import annotations
     from typing import Union, Tuple
     import jax
@@ Expand Down @@