From 324845ebb12831a48eda51d642c6c8829b2e6fb8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= <fabian@schaluck.com>
Date: Thu, 5 Dec 2024 20:19:38 +0000
Subject: [PATCH 1/2] mark jax experimental

---
 python/sdist/amici/jax/__init__.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

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"]

From f1cb7990b5e98ef11def39b4cc48e20ddaef3beb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= <fabian@schaluck.com>
Date: Thu, 5 Dec 2024 20:31:50 +0000
Subject: [PATCH 2/2] add JAX import warning to pyest.ini

---
 pytest.ini | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pytest.ini b/pytest.ini
index adbf313922..8cc45e0fd9 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -12,6 +12,7 @@ filterwarnings =
     ignore:Conservation laws for non-constant species in models with Species-AssignmentRules are currently not supported and will be turned off.:UserWarning
     ignore:Conservation laws for non-constant species in combination with parameterized stoichiometric coefficients are not currently supported and will be turned off.:UserWarning
     ignore:Support for PEtab2.0 is experimental!:UserWarning
+    ignore:The JAX module is experimental and the API may change in the future.:ImportWarning
     # hundreds of SBML <=5.17 warnings
     ignore:.*inspect.getargspec\(\) is deprecated.*:DeprecationWarning
     # pysb warnings