openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/expert/validation.py

@@ -0,0 +1,357 @@
+"""Validation for bring-your-own-functions (byof).
+
+This module provides validation for user-provided JAX functions in expert mode,
+checking signatures, shapes, and differentiability before use.
+"""
+
+import inspect
+from typing import TYPE_CHECKING, List
+
+if TYPE_CHECKING:
+    from openscvx.symbolic.expr.state import State
+
+__all__ = ["validate_byof"]
+
+
+def validate_byof(
+    byof: dict,
+    states: List["State"],
+    n_x: int,
+    n_u: int,
+    N: int = None,
+) -> None:
+    """Validate byof function signatures and shapes.
+
+    Checks that user-provided functions have the correct signatures and return
+    appropriate shapes. Performs validation before functions are used to provide
+    clear error messages.
+
+    Args:
+        byof: Dictionary of user-provided functions to validate
+        states: List of State objects for determining expected shapes
+        n_x: Total dimension of the unified state vector
+        n_u: Total dimension of the unified control vector
+        N: Number of nodes in the trajectory (optional). If provided, validates
+            node indices in nodal constraints.
+
+    Raises:
+        ValueError: If any function has invalid signature or returns wrong shape
+        TypeError: If functions are not callable
+
+    Example:
+        >>> validate_byof(byof, states, n_x=10, n_u=3, N=50)  # Raises if invalid
+    """
+    import jax
+    import jax.numpy as jnp
+
+    # Validate byof keys
+    valid_keys = {"dynamics", "nodal_constraints", "cross_nodal_constraints", "ctcs_constraints"}
+    invalid_keys = set(byof.keys()) - valid_keys
+    if invalid_keys:
+        raise ValueError(f"Unknown byof keys: {invalid_keys}. Valid keys: {valid_keys}")
+
+    # Create dummy inputs for testing
+    dummy_x = jnp.zeros(n_x)
+    dummy_u = jnp.zeros(n_u)
+    dummy_node = 0
+    dummy_params = {}
+
+    # Validate dynamics functions
+    byof_dynamics = byof.get("dynamics", {})
+    if byof_dynamics:
+        # Build mapping from state name to expected shape
+        state_shapes = {state.name: state.shape for state in states}
+
+        for state_name, fn in byof_dynamics.items():
+            if state_name not in state_shapes:
+                raise ValueError(
+                    f"byof dynamics '{state_name}' does not match any state name. "
+                    f"Available states: {list(state_shapes.keys())}"
+                )
+
+            if not callable(fn):
+                raise TypeError(f"byof dynamics '{state_name}' must be callable, got {type(fn)}")
+
+            # Check signature
+            sig = inspect.signature(fn)
+            if len(sig.parameters) != 4:
+                raise ValueError(
+                    f"byof dynamics '{state_name}' must have signature f(x, u, node, params), "
+                    f"got {len(sig.parameters)} parameters: {list(sig.parameters.keys())}"
+                )
+
+            # Test call and check output shape
+            try:
+                result = fn(dummy_x, dummy_u, dummy_node, dummy_params)
+            except Exception as e:
+                raise ValueError(
+                    f"byof dynamics '{state_name}' failed on test call with "
+                    f"x.shape={dummy_x.shape}, u.shape={dummy_u.shape}: {e}"
+                ) from e
+
+            expected_shape = state_shapes[state_name]
+            result_shape = jnp.asarray(result).shape
+            if result_shape != expected_shape:
+                raise ValueError(
+                    f"byof dynamics '{state_name}' returned shape {result_shape}, "
+                    f"expected {expected_shape} (state '{state_name}' shape)"
+                )
+
+            # Test that gradient works (JAX compatibility check)
+            try:
+                jax.grad(lambda x: jnp.sum(fn(x, dummy_u, dummy_node, dummy_params)))(dummy_x)
+            except Exception as e:
+                raise ValueError(
+                    f"byof dynamics '{state_name}' is not differentiable with JAX. "
+                    f"Ensure the function uses JAX operations (jax.numpy, not numpy): {e}"
+                ) from e
+
+    # Validate nodal constraints
+    for i, constraint_spec in enumerate(byof.get("nodal_constraints", [])):
+        if not isinstance(constraint_spec, dict):
+            raise TypeError(
+                f"byof nodal_constraints[{i}] must be a dict (NodalConstraintSpec), "
+                f"got {type(constraint_spec)}"
+            )
+
+        if "constraint_fn" not in constraint_spec:
+            raise ValueError(f"byof nodal_constraints[{i}] missing required key 'constraint_fn'")
+
+        fn = constraint_spec["constraint_fn"]
+        if not callable(fn):
+            raise TypeError(
+                f"byof nodal_constraints[{i}]['constraint_fn'] must be callable, got {type(fn)}"
+            )
+
+        # Check signature
+        sig = inspect.signature(fn)
+        if len(sig.parameters) != 4:
+            raise ValueError(
+                f"byof nodal_constraints[{i}]['constraint_fn'] must have signature "
+                f"f(x, u, node, params), "
+                f"got {len(sig.parameters)} parameters: {list(sig.parameters.keys())}"
+            )
+
+        # Test call
+        try:
+            result = fn(dummy_x, dummy_u, dummy_node, dummy_params)
+        except Exception as e:
+            raise ValueError(
+                f"byof nodal_constraints[{i}]['constraint_fn'] failed on test call with "
+                f"x.shape={dummy_x.shape}, u.shape={dummy_u.shape}: {e}"
+            ) from e
+
+        # Check that result is array-like (can be scalar or vector)
+        try:
+            result_array = jnp.asarray(result)
+        except Exception as e:
+            raise ValueError(
+                f"byof nodal_constraints[{i}]['constraint_fn'] must return array-like value, "
+                f"got {type(result)}: {e}"
+            ) from e
+
+        # Test gradient
+        try:
+            jax.grad(lambda x: jnp.sum(fn(x, dummy_u, dummy_node, dummy_params)))(dummy_x)
+        except Exception as e:
+            raise ValueError(
+                f"byof nodal_constraints[{i}]['constraint_fn'] is not differentiable with JAX: {e}"
+            ) from e
+
+        # Validate nodes if provided
+        if "nodes" in constraint_spec:
+            nodes = constraint_spec["nodes"]
+            if not isinstance(nodes, (list, tuple)):
+                raise TypeError(
+                    f"byof nodal_constraints[{i}]['nodes'] must be a list or tuple, "
+                    f"got {type(nodes)}"
+                )
+            if len(nodes) == 0:
+                raise ValueError(f"byof nodal_constraints[{i}]['nodes'] cannot be empty")
+
+            # Validate node indices if N is provided
+            if N is not None:
+                for node in nodes:
+                    # Handle negative indices (e.g., -1 for last node)
+                    normalized_node = node if node >= 0 else N + node
+                    # Validate range
+                    if not (0 <= normalized_node < N):
+                        raise ValueError(
+                            f"byof nodal_constraints[{i}]['nodes'] contains invalid index {node} "
+                            f"(normalized: {normalized_node}). Valid range is [0, {N}) or "
+                            f"negative indices [-{N}, -1]."
+                        )
+
+    # Validate cross-nodal constraints
+    dummy_X = jnp.zeros((10, n_x))  # Dummy trajectory with 10 nodes
+    dummy_U = jnp.zeros((10, n_u))
+
+    for i, fn in enumerate(byof.get("cross_nodal_constraints", [])):
+        if not callable(fn):
+            raise TypeError(f"byof cross_nodal_constraints[{i}] must be callable, got {type(fn)}")
+
+        # Check signature
+        sig = inspect.signature(fn)
+        if len(sig.parameters) != 3:
+            raise ValueError(
+                f"byof cross_nodal_constraints[{i}] must have signature f(X, U, params), "
+                f"got {len(sig.parameters)} parameters: {list(sig.parameters.keys())}"
+            )
+
+        # Test call
+        try:
+            result = fn(dummy_X, dummy_U, dummy_params)
+        except Exception as e:
+            raise ValueError(
+                f"byof cross_nodal_constraints[{i}] failed on test call with "
+                f"X.shape={dummy_X.shape}, U.shape={dummy_U.shape}: {e}"
+            ) from e
+
+        # Check that result is array-like
+        try:
+            result_array = jnp.asarray(result)
+        except Exception as e:
+            raise ValueError(
+                f"byof cross_nodal_constraints[{i}] must return array-like value, "
+                f"got {type(result)}: {e}"
+            ) from e
+
+        # Test gradient
+        try:
+            jax.grad(lambda X: jnp.sum(fn(X, dummy_U, dummy_params)))(dummy_X)
+        except Exception as e:
+            raise ValueError(
+                f"byof cross_nodal_constraints[{i}] is not differentiable with JAX: {e}"
+            ) from e
+
+    # Validate CTCS constraints
+    for i, ctcs_spec in enumerate(byof.get("ctcs_constraints", [])):
+        if not isinstance(ctcs_spec, dict):
+            raise TypeError(f"byof ctcs_constraints[{i}] must be a dict, got {type(ctcs_spec)}")
+
+        if "constraint_fn" not in ctcs_spec:
+            raise ValueError(f"byof ctcs_constraints[{i}] missing required key 'constraint_fn'")
+
+        fn = ctcs_spec["constraint_fn"]
+        if not callable(fn):
+            raise TypeError(
+                f"byof ctcs_constraints[{i}]['constraint_fn'] must be callable, got {type(fn)}"
+            )
+
+        # Check signature
+        sig = inspect.signature(fn)
+        if len(sig.parameters) != 4:
+            raise ValueError(
+                f"byof ctcs_constraints[{i}]['constraint_fn'] must have signature "
+                f"f(x, u, node, params), got {len(sig.parameters)} parameters: "
+                f"{list(sig.parameters.keys())}"
+            )
+
+        # Test call
+        try:
+            result = fn(dummy_x, dummy_u, dummy_node, dummy_params)
+        except Exception as e:
+            raise ValueError(
+                f"byof ctcs_constraints[{i}]['constraint_fn'] failed on test call: {e}"
+            ) from e
+
+        # Check that result is scalar
+        result_array = jnp.asarray(result)
+        if result_array.shape != ():
+            raise ValueError(
+                f"byof ctcs_constraints[{i}]['constraint_fn'] must return a scalar, "
+                f"got shape {result_array.shape}"
+            )
+
+        # Test gradient
+        try:
+            jax.grad(lambda x: fn(x, dummy_u, dummy_node, dummy_params))(dummy_x)
+        except Exception as e:
+            raise ValueError(
+                f"byof ctcs_constraints[{i}]['constraint_fn'] is not differentiable with JAX: {e}"
+            ) from e
+
+        # Validate penalty function if provided
+        if "penalty" in ctcs_spec:
+            penalty_spec = ctcs_spec["penalty"]
+            if callable(penalty_spec):
+                # Test custom penalty function
+                try:
+                    test_residual = jnp.array(0.5)
+                    penalty_result = penalty_spec(test_residual)
+                    jnp.asarray(penalty_result)
+                except Exception as e:
+                    raise ValueError(
+                        f"byof ctcs_constraints[{i}]['penalty'] custom function failed: {e}"
+                    ) from e
+            elif penalty_spec not in ["square", "l1", "huber"]:
+                raise ValueError(
+                    f"byof ctcs_constraints[{i}]['penalty'] must be 'square', 'l1', 'huber', "
+                    f"or a callable, got {penalty_spec!r}"
+                )
+
+        # Validate idx if provided
+        if "idx" in ctcs_spec:
+            idx = ctcs_spec["idx"]
+            if not isinstance(idx, int):
+                raise TypeError(
+                    f"byof ctcs_constraints[{i}]['idx'] must be an integer, got {type(idx)}"
+                )
+            if idx < 0:
+                raise ValueError(
+                    f"byof ctcs_constraints[{i}]['idx'] must be non-negative, got {idx}"
+                )
+
+        # Validate bounds if provided
+        if "bounds" in ctcs_spec:
+            bounds = ctcs_spec["bounds"]
+            if not isinstance(bounds, (tuple, list)) or len(bounds) != 2:
+                raise ValueError(
+                    f"byof ctcs_constraints[{i}]['bounds'] must be a (min, max) tuple, got {bounds}"
+                )
+            if bounds[0] > bounds[1]:
+                raise ValueError(
+                    f"byof ctcs_constraints[{i}]['bounds'] min ({bounds[0]}) must be <= "
+                    f"max ({bounds[1]})"
+                )
+        else:
+            # Use default bounds for initial value validation
+            bounds = (0.0, 1e-4)
+
+        # Validate initial value is within bounds
+        if "initial" in ctcs_spec:
+            initial = ctcs_spec["initial"]
+            if not (bounds[0] <= initial <= bounds[1]):
+                raise ValueError(
+                    f"byof ctcs_constraints[{i}]['initial'] ({initial}) must be within "
+                    f"bounds [{bounds[0]}, {bounds[1]}]"
+                )
+
+        # Validate over (node interval) if provided
+        if "over" in ctcs_spec:
+            over = ctcs_spec["over"]
+            if not isinstance(over, (tuple, list)) or len(over) != 2:
+                raise ValueError(
+                    f"byof ctcs_constraints[{i}]['over'] must be a (start, end) tuple, got {over}"
+                )
+            start, end = over
+            if not isinstance(start, int) or not isinstance(end, int):
+                raise TypeError(
+                    f"byof ctcs_constraints[{i}]['over'] indices must be integers, "
+                    f"got start={type(start)}, end={type(end)}"
+                )
+            if start >= end:
+                raise ValueError(
+                    f"byof ctcs_constraints[{i}]['over'] start ({start}) must be < end ({end})"
+                )
+            if start < 0:
+                raise ValueError(
+                    f"byof ctcs_constraints[{i}]['over'] start ({start}) must be non-negative"
+                )
+            # Validate against trajectory length if N is provided
+            if N is not None:
+                if end > N:
+                    raise ValueError(
+                        f"byof ctcs_constraints[{i}]['over'] end ({end}) exceeds "
+                        f"trajectory length ({N})"
+                    )

openscvx/integrators/__init__.py

@@ -0,0 +1,48 @@
+"""Numerical integration schemes for trajectory optimization.
+
+This module provides implementations of numerical integrators used for simulating
+continuous-time dynamics.
+
+Current Implementations:
+    RK45 Integration: Explicit Runge-Kutta-Fehlberg method (4th/5th order)
+        with both fixed-step and adaptive implementations via Diffrax.
+        Supports a variety of explicit and implicit ODE solvers through the
+        Diffrax backend (Dopri5/8, Tsit5, KenCarp3/4/5, etc.).
+
+Planned Architecture (ABC-based):
+
+A base class will be introduced to enable pluggable integrator implementations.
+This will enable users to implement custom integrators.
+Future integrators will implement the Integrator interface:
+
+```python
+# integrators/base.py (planned):
+class Integrator(ABC):
+    @abstractmethod
+    def step(self, f: Callable, x: Array, u: Array, t: float, dt: float) -> Array:
+        '''Take one integration step from state x at time t with step dt.'''
+        ...
+
+    @abstractmethod
+    def integrate(self, f: Callable, x0: Array, u_traj: Array,
+                    t_span: tuple[float, float], num_steps: int) -> Array:
+        '''Integrate over a time span with given control trajectory.'''
+        ...
+```
+"""
+
+from .runge_kutta import (
+    SOLVER_MAP,
+    rk45_step,
+    solve_ivp_diffrax,
+    solve_ivp_diffrax_prop,
+    solve_ivp_rk45,
+)
+
+__all__ = [
+    "SOLVER_MAP",
+    "rk45_step",
+    "solve_ivp_rk45",
+    "solve_ivp_diffrax",
+    "solve_ivp_diffrax_prop",
+]

openscvx/integrators/runge_kutta.py

@@ -0,0 +1,281 @@
+import os
+from typing import Any, Callable
+
+import diffrax as dfx
+import jax
+import jax.numpy as jnp
+from diffrax._global_interpolation import DenseInterpolation
+from jax import tree_util
+
+os.environ["EQX_ON_ERROR"] = "nan"
+
+
+# Safely check if DenseInterpolation is already registered
+try:
+    # Attempt to flatten a dummy DenseInterpolation instance
+    # Provide dummy arguments to create a valid instance
+    dummy_instance = DenseInterpolation(
+        ts=jnp.array([]),
+        ts_size=0,
+        infos=None,
+        interpolation_cls=None,
+        direction=None,
+        t0_if_trivial=0.0,
+        y0_if_trivial=jnp.array([]),
+    )
+    tree_util.tree_flatten(dummy_instance)
+except ValueError:
+    # Register DenseInterpolation as a PyTree node if not already registered
+    def dense_interpolation_flatten(obj):
+        # Flatten the internal data of DenseInterpolation
+        return (obj._data,), None
+
+    def dense_interpolation_unflatten(aux_data, children):
+        # Reconstruct DenseInterpolation from its flattened data
+        return DenseInterpolation(*children)
+
+    tree_util.register_pytree_node(
+        DenseInterpolation,
+        dense_interpolation_flatten,
+        dense_interpolation_unflatten,
+    )
+
+SOLVER_MAP = {
+    "Tsit5": dfx.Tsit5,
+    "Euler": dfx.Euler,
+    "Heun": dfx.Heun,
+    "Midpoint": dfx.Midpoint,
+    "Ralston": dfx.Ralston,
+    "Dopri5": dfx.Dopri5,
+    "Dopri8": dfx.Dopri8,
+    "Bosh3": dfx.Bosh3,
+    "ReversibleHeun": dfx.ReversibleHeun,
+    "ImplicitEuler": dfx.ImplicitEuler,
+    "KenCarp3": dfx.KenCarp3,
+    "KenCarp4": dfx.KenCarp4,
+    "KenCarp5": dfx.KenCarp5,
+}
+
+
+# fmt: off
+def rk45_step(
+    f: Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray],
+    t: jnp.ndarray,
+    y: jnp.ndarray,
+    h: float,
+    *args
+) -> jnp.ndarray:
+    """
+    Perform a single RK45 (Runge-Kutta-Fehlberg) integration step.
+
+    This implements the classic Dorman-Prince coefficients for an
+    explicit 4(5) method, returning the fourth-order estimate.
+
+    Args:
+        f (Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray]):
+            ODE right-hand side; signature f(t, y, *args) -> dy/dt.
+        t (jnp.ndarray): Current time.
+        y (jnp.ndarray): Current state vector.
+        h (float): Step size.
+        *args: Additional arguments passed to `f`.
+
+    Returns:
+        jnp.ndarray: Next state estimate at t + h.
+    """
+    k1 = f(t, y, *args)
+    k2 = f(t + h/4, y + h*k1/4, *args)
+    k3 = f(t + 3*h/8, y + 3*h*k1/32 + 9*h*k2/32, *args)
+    k4 = f(t + 12*h/13, y + 1932*h*k1/2197 - 7200*h*k2/2197 + 7296*h*k3/2197, *args)
+    k5 = f(t + h, y + 439*h*k1/216 - 8*h*k2 + 3680*h*k3/513 - 845*h*k4/4104, *args)
+    y_next = y + h * (25*k1/216 + 1408*k3/2565 + 2197*k4/4104 - k5/5)
+    return y_next
+# fmt: on
+
+
+def solve_ivp_rk45(
+    f: Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray],
+    tau_final: float,
+    y_0: jnp.ndarray,
+    args,
+    tau_0: float = 0.0,
+    num_substeps: int = 50,
+    is_not_compiled: bool = False,
+):
+    """
+    Solve an initial-value ODE problem using fixed-step RK45 integration.
+
+    Args:
+        f (Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray]):
+            ODE right-hand side; signature f(t, y, *args) -> dy/dt.
+        tau_final (float): Final integration time.
+        y_0 (jnp.ndarray): Initial state at tau_0.
+        args (tuple): Extra arguments to pass to `f`.
+        tau_0 (float, optional): Initial time. Defaults to 0.0.
+        num_substeps (int, optional): Number of output time points. Defaults to 50.
+        is_not_compiled (bool, optional): If True, use Python loop instead of
+            JAX `lax.fori_loop`. Defaults to False.
+
+    Returns:
+        jnp.ndarray: Array of shape (num_substeps, state_dim) with solution at each time.
+    """
+    substeps = jnp.linspace(tau_0, tau_final, num_substeps)
+
+    h = (tau_final - tau_0) / (len(substeps) - 1)
+    solution = jnp.zeros((len(substeps), len(y_0)))
+    solution = solution.at[0].set(y_0)
+
+    if is_not_compiled:
+        for i in range(1, len(substeps)):
+            t = tau_0 + i * h
+            solution = solution.at[i].set(rk45_step(f, t, solution[i - 1], h, *args))
+    else:
+
+        def body_fun(i, val):
+            t, y, V_result = val
+            y_next = rk45_step(f, t, y, h, *args)
+            V_result = V_result.at[i].set(y_next)
+            return (t + h, y_next, V_result)
+
+        _, _, solution = jax.lax.fori_loop(1, len(substeps), body_fun, (tau_0, y_0, solution))
+
+    return solution
+
+
+def solve_ivp_diffrax(
+    f: Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray],
+    tau_final: float,
+    y_0: jnp.ndarray,
+    args,
+    tau_0: float = 0.0,
+    num_substeps: int = 50,
+    solver_name: str = "Dopri8",
+    rtol: float = 1e-3,
+    atol: float = 1e-6,
+    extra_kwargs=None,
+):
+    """
+    Solve an initial-value ODE problem using a Diffrax adaptive solver.
+
+    Args:
+        f (Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray]):
+            ODE right-hand side; f(t, y, *args).
+        tau_final (float): Final integration time.
+        y_0 (jnp.ndarray): Initial state at tau_0.
+        args (tuple): Extra arguments to pass to `f` in the solver term.
+        tau_0 (float, optional): Initial time. Defaults to 0.0.
+        num_substeps (int, optional): Number of save points between tau_0 and tau_final.
+            Defaults to 50.
+        solver_name (str, optional): Key into SOLVER_MAP for the Diffrax solver class.
+            Defaults to "Dopri8".
+        rtol (float, optional): Relative tolerance for adaptive stepping. Defaults to 1e-3.
+        atol (float, optional): Absolute tolerance for adaptive stepping. Defaults to 1e-6.
+        extra_kwargs (dict, optional): Additional keyword arguments forwarded to `diffeqsolve`.
+
+    Returns:
+        jnp.ndarray: Solution states at the requested save points, shape (num_substeps, state_dim).
+
+    Raises:
+        ValueError: If `solver_name` is not in SOLVER_MAP.
+    """
+    substeps = jnp.linspace(tau_0, tau_final, num_substeps)
+
+    solver_class = SOLVER_MAP.get(solver_name)
+    if solver_class is None:
+        raise ValueError(f"Unknown solver: {solver_name}")
+    solver = solver_class()
+
+    term = dfx.ODETerm(lambda t, y, args: f(t, y, *args))
+    stepsize_controller = dfx.PIDController(rtol=rtol, atol=atol)
+    solution = dfx.diffeqsolve(
+        term,
+        solver=solver,
+        t0=tau_0,
+        t1=tau_final,
+        dt0=(tau_final - tau_0) / (len(substeps) - 1),
+        y0=y_0,
+        args=args,
+        stepsize_controller=stepsize_controller,
+        saveat=dfx.SaveAt(ts=substeps),
+        progress_meter=dfx.NoProgressMeter(),
+        **(extra_kwargs or {}),
+    )
+
+    return solution.ys
+
+
+def solve_ivp_diffrax_prop(
+    f: Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray],
+    tau_final: float,
+    y_0: jnp.ndarray,
+    args,
+    tau_0: float = 0.0,
+    num_substeps: int = 50,
+    solver_name: str = "Dopri8",
+    rtol: float = 1e-3,
+    atol: float = 1e-6,
+    extra_kwargs=None,
+    save_time: jnp.ndarray = None,
+    mask: jnp.ndarray = None,
+):
+    """
+    Solve an initial-value ODE problem using a Diffrax adaptive solver.
+    This function is specifically designed for use in the context of
+    trajectory optimization and handles the nonlinear single-shot propagation
+    of state variables in undilated time.
+
+    Args:
+        f (Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray]): ODE right-hand side;
+            signature f(t, y, *args) -> dy/dt.
+        tau_final (float): Final integration time.
+        y_0 (jnp.ndarray): Initial state at tau_0.
+        args (tuple): Extra arguments to pass to `f` in the solver term.
+        tau_0 (float, optional): Initial time. Defaults to 0.0.
+        num_substeps (int, optional): Number of save points between tau_0 and tau_final.
+            Defaults to 50.
+        solver_name (str, optional): Key into SOLVER_MAP for the Diffrax solver class.
+            Defaults to "Dopri8".
+        rtol (float, optional): Relative tolerance for adaptive stepping. Defaults to 1e-3.
+        atol (float, optional): Absolute tolerance for adaptive stepping. Defaults to 1e-6.
+        extra_kwargs (dict, optional): Additional keyword arguments forwarded to `diffeqsolve`.
+        save_time (jnp.ndarray, optional): Time points at which to evaluate the solution.
+            Must be provided for export compatibility.
+        mask (jnp.ndarray, optional): Boolean mask for the save_time points.
+
+    Returns:
+        jnp.ndarray: Solution states at the requested save points, shape (num_substeps, state_dim).
+    Raises:
+        ValueError: If `solver_name` is not in SOLVER_MAP or if save_time is not provided.
+    """
+
+    if save_time is None:
+        raise ValueError("save_time must be provided for export compatibility.")
+    if mask is None:
+        mask = jnp.ones_like(save_time, dtype=bool)
+
+    solver_class = SOLVER_MAP.get(solver_name)
+    if solver_class is None:
+        raise ValueError(f"Unknown solver: {solver_name}")
+    solver = solver_class()
+
+    term = dfx.ODETerm(lambda t, y, args: f(t, y, *args))
+    stepsize_controller = dfx.PIDController(rtol=rtol, atol=atol)
+
+    solution = dfx.diffeqsolve(
+        term,
+        solver=solver,
+        t0=tau_0,
+        t1=tau_final,
+        dt0=(tau_final - tau_0) / 1,
+        y0=y_0,
+        args=args,
+        stepsize_controller=stepsize_controller,
+        saveat=dfx.SaveAt(dense=True),
+        **(extra_kwargs or {}),
+    )
+
+    # Evaluate all save_time points (static size), then mask them
+    all_evals = jax.vmap(solution.evaluate)(save_time)  # shape: (MAX_TAU_LEN, n_states)
+    masked_array = jnp.where(mask[:, None], all_evals, jnp.zeros_like(all_evals))
+    # shape: (variable_len, n_states)
+
+    return masked_array