fixes & remove sciml dependency

python/sdist/amici/jax/nn.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 
-from petab_sciml import MLModel, Layer, Node
+
 import equinox as eqx
 import jax.numpy as jnp
 
@@ -30,7 +30,10 @@ def tanhshrink(x: jnp.ndarray) -> jnp.ndarray:
     return x - jnp.tanh(x)
 
 
-def generate_equinox(ml_model: MLModel, filename: Path | str):
+def generate_equinox(ml_model: "MLModel", filename: Path | str):  # noqa: F821
+    # TODO: move to top level import and replace forward type definitions
+    from petab_sciml import Layer
+
     filename = Path(filename)
     layer_indent = 12
     node_indent = 8
@@ -87,7 +90,7 @@ def _process_argval(v):
     return str(v)
 
 
-def _generate_layer(layer: Layer, indent: int, ilayer: int) -> str:
+def _generate_layer(layer: "Layer", indent: int, ilayer: int) -> str:  # noqa: F821
     layer_map = {
         "Dropout1d": "eqx.nn.Dropout",
         "Dropout2d": "eqx.nn.Dropout",
@@ -146,7 +149,7 @@ def _generate_layer(layer: Layer, indent: int, ilayer: int) -> str:
     return f"{' ' * indent}'{layer.layer_id}': {layer_str}"
 
 
-def _generate_forward(node: Node, indent, layer_type=str) -> str:
+def _generate_forward(node: "Node", indent, layer_type=str) -> str:  # noqa: F821
     if node.op == "placeholder":
         # TODO: inconsistent target vs name
         return f"{' ' * indent}{node.name} = input"

python/sdist/amici/jax/ode_export.py

@@ -130,7 +130,7 @@ def __init__(
         outdir: Path | str | None = None,
         verbose: bool | int | None = False,
         model_name: str | None = "model",
-        hybridisation: dict[str, str] = {},
+        hybridisation: dict[str, dict] = None,
     ):
         """
         Generate AMICI jax files for the ODE provided to the constructor.
@@ -159,7 +159,7 @@ def __init__(
 
         self.model: DEModel = ode_model
 
-        self.hybridisation = hybridisation
+        self.hybridisation = hybridisation if hybridisation is not None else {}
 
         self._code_printer = AmiciJaxCodePrinter()
 

tests/sciml/test_sciml.py

@@ -327,8 +327,9 @@ def test_ude(test):
     # gradient
     sllh, _ = eqx.filter_grad(run_simulations, has_aux=True)(
         jax_problem,
-        solver=diffrax.Tsit5(),
-        controller=diffrax.PIDController(atol=1e-10, rtol=1e-10),
+        solver=diffrax.Kvaerno5(),
+        controller=diffrax.PIDController(atol=1e-14, rtol=1e-14),
+        max_steps=2**16,
     )
     expected = (
         pd.concat(