diff --git a/python/sdist/amici/jax/__init__.py b/python/sdist/amici/jax/__init__.py
index e14d231e1e..8b67abda27 100644
--- a/python/sdist/amici/jax/__init__.py
+++ b/python/sdist/amici/jax/__init__.py
@@ -1,6 +1,21 @@
-"""Interface to facilitate AMICI generated models using JAX"""
+"""
+JAX
+---
+
+This module provides an interface to generate and use AMICI models with JAX. Please note that this module is
+experimental, the API may substantially change in the future. Use at your own risk and do not expect backward
+compatibility.
+"""
+
+from warnings import warn
 
 from amici.jax.petab import JAXProblem, run_simulations
 from amici.jax.model import JAXModel
 
+warn(
+    "The JAX module is experimental and the API may change in the future.",
+    ImportWarning,
+    stacklevel=2,
+)
+
 __all__ = ["JAXModel", "JAXProblem", "run_simulations"]