Fix bugs in Python 3.9

chaobrain · Jul 13, 2024 · 2c9fe98 · 2c9fe98
1 parent f95d5f9
commit 2c9fe98
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/brainunit/__init__.py b/brainunit/__init__.py
@@ -15,6 +15,7 @@
 
 __version__ = "0.0.1.1"
 
+from __future__ import annotations
 import importlib
 
 from . import math

diff --git a/brainunit/_base.py b/brainunit/_base.py
@@ -18,15 +18,14 @@
 import collections
 import itertools
 import numbers
-from functools import wraps
 import operator
 from contextlib import contextmanager
+from functools import wraps
 from typing import Union, Optional, Sequence, Callable, Tuple, Any, List
 
 import jax
 import jax.numpy as jnp
 import numpy as np
-from jax import Array
 from jax.interpreters.partial_eval import DynamicJaxprTracer
 from jax.tree_util import register_pytree_node_class