diff --git a/dendritex/_base.py b/dendritex/_base.py
index c34712a..ff108dd 100644
--- a/dendritex/_base.py
+++ b/dendritex/_base.py
@@ -16,6 +16,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 from typing import Optional, Dict, Sequence, Callable, NamedTuple, Tuple
 
 import brainstate as bst
diff --git a/dendritex/_base_test.py b/dendritex/_base_test.py
index bed3fa6..fa6d048 100644
--- a/dendritex/_base_test.py
+++ b/dendritex/_base_test.py
@@ -15,7 +15,7 @@
 
 import unittest
 
+
 class TestIon(unittest.TestCase):
   def test1(self):
     pass
-
diff --git a/dendritex/_integrators.py b/dendritex/_integrators.py
index 08ec793..0318b1d 100644
--- a/dendritex/_integrators.py
+++ b/dendritex/_integrators.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
+
 import brainstate as bst
 import brainunit as bu
 import jax
@@ -37,7 +39,6 @@ def tree_map(f, tree, *rest):
   return jax.tree.map(f, tree, *rest, is_leaf=lambda a: isinstance(a, bu.Quantity))
 
 
-
 def euler_step(target, t: jax.typing.ArrayLike, *args):
   with bst.environ.context(t=t):
     return target(*args)
diff --git a/dendritex/channels/calcium.py b/dendritex/channels/calcium.py
index 2b5c6b5..794473b 100644
--- a/dendritex/channels/calcium.py
+++ b/dendritex/channels/calcium.py
@@ -5,6 +5,8 @@
 
 """
 
+from __future__ import annotations
+
 from typing import Union, Callable, Optional
 
 import brainstate as bst
diff --git a/dendritex/channels/hyperpolarization_activated.py b/dendritex/channels/hyperpolarization_activated.py
index 216f7ac..02e3cdf 100644
--- a/dendritex/channels/hyperpolarization_activated.py
+++ b/dendritex/channels/hyperpolarization_activated.py
@@ -4,6 +4,8 @@
 This module implements hyperpolarization-activated cation channels.
 """
 
+from __future__ import annotations
+
 from typing import Union, Callable, Optional
 
 import brainstate as bst
diff --git a/dendritex/channels/leaky.py b/dendritex/channels/leaky.py
index 0882830..8a03bb6 100644
--- a/dendritex/channels/leaky.py
+++ b/dendritex/channels/leaky.py
@@ -5,6 +5,8 @@
 
 """
 
+from __future__ import annotations
+
 from typing import Union, Callable, Sequence, Optional
 
 import brainstate as bst
diff --git a/dendritex/channels/potassium.py b/dendritex/channels/potassium.py
index 6529a88..5937d36 100644
--- a/dendritex/channels/potassium.py
+++ b/dendritex/channels/potassium.py
@@ -5,6 +5,8 @@
 
 """
 
+from __future__ import annotations
+
 from typing import Union, Callable, Optional, Sequence
 
 import brainstate as bst
diff --git a/dendritex/channels/potassium_calcium.py b/dendritex/channels/potassium_calcium.py
index cb24682..a1bd869 100644
--- a/dendritex/channels/potassium_calcium.py
+++ b/dendritex/channels/potassium_calcium.py
@@ -5,6 +5,8 @@
 This module implements calcium-dependent potassium channels.
 """
 
+from __future__ import annotations
+
 from typing import Union, Callable, Optional
 
 import brainstate as bst
diff --git a/dendritex/channels/sodium.py b/dendritex/channels/sodium.py
index cab1ba5..98ca964 100644
--- a/dendritex/channels/sodium.py
+++ b/dendritex/channels/sodium.py
@@ -5,6 +5,8 @@
 
 """
 
+from __future__ import annotations
+
 from typing import Union, Callable, Optional
 
 import brainstate as bst
diff --git a/dendritex/ions/calcium.py b/dendritex/ions/calcium.py
index 95d1eb0..074ccce 100644
--- a/dendritex/ions/calcium.py
+++ b/dendritex/ions/calcium.py
@@ -15,6 +15,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 from typing import Union, Callable, Optional
 
 import brainstate as bst
diff --git a/dendritex/ions/potassium.py b/dendritex/ions/potassium.py
index 5968c3d..5c7fe7d 100644
--- a/dendritex/ions/potassium.py
+++ b/dendritex/ions/potassium.py
@@ -14,10 +14,12 @@
 # ==============================================================================
 
 
+from __future__ import annotations
+
 from typing import Union, Callable, Optional
 
-import brainunit as bu
 import brainstate as bst
+import brainunit as bu
 
 from .._base import Ion, Channel, check_hierarchies
 
diff --git a/dendritex/ions/sodium.py b/dendritex/ions/sodium.py
index 4c45c9f..73ce7e4 100644
--- a/dendritex/ions/sodium.py
+++ b/dendritex/ions/sodium.py
@@ -13,13 +13,14 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
 
 from typing import Union, Callable, Optional
 
 import brainstate as bst
 import brainunit as bu
 
-from .._base import Ion, Channel, check_hierarchies
+from .._base import Ion
 
 __all__ = [
   'Sodium',
@@ -57,4 +58,3 @@ def __init__(
     )
     self.E = bst.init.param(E, self.varshape, allow_none=False)
     self.C = bst.init.param(C, self.varshape, allow_none=False)
-
diff --git a/dendritex/neurons/__init__.py b/dendritex/neurons/__init__.py
index 2d24f77..5891c52 100644
--- a/dendritex/neurons/__init__.py
+++ b/dendritex/neurons/__init__.py
@@ -17,6 +17,4 @@
 from .hh import *
 from .hh import __all__ as _hh_all
 
-
 __all__ = _hh_all
-
diff --git a/dendritex/neurons/hh.py b/dendritex/neurons/hh.py
index df1f358..195cf88 100644
--- a/dendritex/neurons/hh.py
+++ b/dendritex/neurons/hh.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
 
 from typing import Union, Optional, Callable