diff --git a/src/diffraxtra/__init__.py b/src/diffraxtra/__init__.py index 1644568..9391555 100644 --- a/src/diffraxtra/__init__.py +++ b/src/diffraxtra/__init__.py @@ -6,9 +6,12 @@ "VectorizedDenseInterpolation", ] +from ._src import ( + AbstractVectorizedDenseInterpolation, + DiffEqSolver, + VectorizedDenseInterpolation, +) from ._version import ( # noqa: F401 version as __version__, version_tuple as __version_tuple__, ) -from .diffeq import DiffEqSolver -from .interp import AbstractVectorizedDenseInterpolation, VectorizedDenseInterpolation diff --git a/src/diffraxtra/_src/__init__.py b/src/diffraxtra/_src/__init__.py new file mode 100644 index 0000000..904036f --- /dev/null +++ b/src/diffraxtra/_src/__init__.py @@ -0,0 +1,13 @@ +"""Extras for `diffrax`. Private API.""" + +__all__ = [ + "DiffEqSolver", + "AbstractVectorizedDenseInterpolation", + "VectorizedDenseInterpolation", +] + +from .diffeq import DiffEqSolver +from .interp import ( + AbstractVectorizedDenseInterpolation, + VectorizedDenseInterpolation, +) diff --git a/src/diffraxtra/diffeq.py b/src/diffraxtra/_src/diffeq.py similarity index 100% rename from src/diffraxtra/diffeq.py rename to src/diffraxtra/_src/diffeq.py diff --git a/src/diffraxtra/interp.py b/src/diffraxtra/_src/interp.py similarity index 100% rename from src/diffraxtra/interp.py rename to src/diffraxtra/_src/interp.py