From 7a780ad100f97938ed4e09d2972115c9a9097079 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= <fabian@schaluck.com>
Date: Thu, 5 Dec 2024 20:01:49 +0000
Subject: [PATCH] fix doc typehints

---
 python/sdist/amici/jax/petab.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/sdist/amici/jax/petab.py b/python/sdist/amici/jax/petab.py
index b9f4b18881..bb8749e27c 100644
--- a/python/sdist/amici/jax/petab.py
+++ b/python/sdist/amici/jax/petab.py
@@ -154,7 +154,7 @@ def _get_parameter_mappings(
     def _get_measurements(
         self, simulation_conditions: pd.DataFrame
     ) -> dict[
-        str,
+        tuple[str, ...],
         tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray],
     ]:
         """
@@ -320,7 +320,7 @@ def run_simulation(
             Step size controller to use for simulation
         :param max_steps:
             Maximum number of steps to take during simulation
-        :param preeq:
+        :param x_preeq:
             Pre-equilibration state if available
         :return:
             Tuple of log-likelihood and simulation statistics
@@ -375,7 +375,7 @@ def run_preequilibration(
 
 def run_simulations(
     problem: JAXProblem,
-    simulation_conditions: Iterable[tuple] | None = None,
+    simulation_conditions: Iterable[tuple[str, ...]] | None = None,
     solver: diffrax.AbstractSolver = diffrax.Kvaerno5(),
     controller: diffrax.AbstractStepSizeController = diffrax.PIDController(
         rtol=1e-8,