diff --git a/src/pmhn/_trees/__init__.py b/src/pmhn/_trees/__init__.py
index daac756..744a8e9 100644
--- a/src/pmhn/_trees/__init__.py
+++ b/src/pmhn/_trees/__init__.py
@@ -7,9 +7,12 @@
 https://doi.org/10.1038/s41467-023-39400-w
 """
 
-from pmhn._trees._simulate import simulate_trees
-from pmhn._trees._interfaces import Tree
+# TODO(Pawel): test if this implementation works.
+
+
 from pmhn._trees._backend import OriginalTreeMHNBackend
+from pmhn._trees._interfaces import Tree
+from pmhn._trees._simulate import simulate_trees
 
 __all__ = [
     "simulate_trees",