cgarciae · cgarciae · Mar 17, 2023 · Mar 17, 2023
diff --git a/simple_pytree/__init__.py b/simple_pytree/__init__.py
@@ -1,12 +1,3 @@
 __version__ = "0.1.5"
 
 from .pytree import Pytree, field, static_field
-from .ref import (
-    PytreeRef,
-    Ref,
-    RefField,
-    clone_references,
-    cross_barrier,
-    incremented_ref,
-    ref_field,
-)
diff --git a/simple_pytree/ids.py b/simple_pytree/ids.py
diff --git a/simple_pytree/ref.py b/simple_pytree/ref.py
diff --git a/simple_pytree/tracers.py b/simple_pytree/tracers.py
diff --git a/tests/test_pytree.py b/tests/test_pytree.py
@@ -134,7 +134,7 @@ def __init__(self, y) -> None:
                 self.x = 2
                 self.y = y
 
-        pytree: Foo = Foo(y=3)
+        pytree = Foo(y=3)
 
         leaves = jax.tree_util.tree_leaves(pytree)
         assert leaves == [3]