diff --git a/python/sdist/amici/__init__.py b/python/sdist/amici/__init__.py
index 6788fefe77..0a7d3c6581 100644
--- a/python/sdist/amici/__init__.py
+++ b/python/sdist/amici/__init__.py
@@ -141,8 +141,6 @@ def get_model(self) -> amici.Model:
             """Create a model instance."""
             ...
 
-        def get_jax_model(self) -> JAXModel: ...
-
     AmiciModel = Union[amici.Model, amici.ModelPtr]
 else:
     ModelModule = ModuleType
diff --git a/python/sdist/amici/__init__.template.py b/python/sdist/amici/__init__.template.py
index abd07a81ab..06302eba9d 100644
--- a/python/sdist/amici/__init__.template.py
+++ b/python/sdist/amici/__init__.template.py
@@ -2,13 +2,9 @@
 
 import sys
 from pathlib import Path
-from typing import TYPE_CHECKING
 import amici
 
 
-if TYPE_CHECKING:
-    pass
-
 # Ensure we are binary-compatible, see #556
 if "TPL_AMICI_VERSION" != amici.__version__:
     raise amici.AmiciVersionError(