diff --git a/src/pystencils_autodiff/__init__.py b/src/pystencils_autodiff/__init__.py
index fafab5ab22e7ac018d4feeac67ef7dd49f339c4a..4a708898f94d9905e8a79004de5353433e3087bc 100644
--- a/src/pystencils_autodiff/__init__.py
+++ b/src/pystencils_autodiff/__init__.py
@@ -6,7 +6,7 @@ from pystencils_autodiff._field_to_tensors import (  # NOQA
     tf_variable_from_field, torch_tensor_from_field)
 from pystencils_autodiff._adjoint_field import AdjointField
 from pystencils_autodiff._autodiff import (
-    AutoDiffAstPair, AutoDiffOp, create_backward_assignments, get_jacobian_of_assignments)
+    AutoDiffAstPair, AutoDiffOp, create_backward_assignments, get_jacobian_of_assignments, DiffModes)
 
 __all__ = ['backends',
            'AdjointField',
@@ -16,7 +16,7 @@ __all__ = ['backends',
            'AutoDiffAstPair',
            "tf_constant_from_field", " tf_placeholder_from_field",
            "tf_scalar_variable_from_field", " tf_variable_from_field",
-           "torch_tensor_from_field"]
+           "torch_tensor_from_field", "DiffModes"]
 
 sys.modules['pystencils.autodiff'] = pystencils_autodiff
 sys.modules['pystencils.autodiff.backends'] = pystencils_autodiff.backends
diff --git a/tests/test_autodiff.py b/tests/test_autodiff.py
index 3ddc3a791aab283a7d3fe582f888fbecf2bbf263..8d3f5354f327c20a1f3ef3a346250bebf18e6cf8 100644
--- a/tests/test_autodiff.py
+++ b/tests/test_autodiff.py
@@ -2,7 +2,8 @@ import sympy as sp
 
 import pystencils as ps
 import pystencils_autodiff
-from pystencils_autodiff.autodiff import DiffModes
+from pystencils_autodiff import DiffModes
+
 
 
 def test_simple_2d_check_assignment_collection():