diff --git a/dynestyx/inference/configs/filter.py b/dynestyx/inference/configs/filter.py
index 1590d935..d3d18d19 100644
--- a/dynestyx/inference/configs/filter.py
+++ b/dynestyx/inference/configs/filter.py
@@ -5,8 +5,8 @@
 import math
 from typing import Literal
 
-import jax
 import jax.random as jr
+from jaxtyping import PRNGKeyArray
 
 ResamplingBaseMethod = Literal["systematic", "multinomial", "stratified"]
 ResamplingDifferentiableMethod = Literal["stop_gradient", "straight_through", "soft"]
@@ -55,7 +55,7 @@ class BaseFilterConfig(abc.ABC):
             this factor before the update. Values slightly above `1.0`
             implement covariance inflation, which can improve robustness when
             the model is misspecified. `None` disables rescaling.
-        crn_seed (jax.Array | None): Fix the PRNG key for stochastic filters
+        crn_seed (PRNGKeyArray | None): Fix the PRNG key for stochastic filters
             (EnKF, PF). Useful when differentiating through the filter:
             a fixed key makes the randomness a deterministic function of model
             parameters. `None` draws a fresh key each call.
@@ -78,7 +78,7 @@ class BaseFilterConfig(abc.ABC):
     record_max_elems: int = 100_000
     filter_source: FilterSource | None = None
     cov_rescaling: float | None = None
-    crn_seed: jax.Array | None = None
+    crn_seed: PRNGKeyArray | None = None
 
 
 @dataclasses.dataclass
@@ -107,7 +107,7 @@ class EnKFConfig(BaseFilterConfig):
         n_particles (int): Number of ensemble members. More members give a
             better covariance estimate at higher compute cost. Defaults to
             `30`.
-        crn_seed (jax.Array | None): Fixed PRNG key for the ensemble. Defaults
+        crn_seed (PRNGKeyArray | None): Fixed PRNG key for the ensemble. Defaults
             to `jr.PRNGKey(0)`, i.e., common random numbers are used. This
             can reduce variance in gradient-based learning, but introduces
             further bias.
@@ -163,7 +163,7 @@ class EnKFConfig(BaseFilterConfig):
     """
 
     n_particles: int = 30
-    crn_seed: jax.Array | None = dataclasses.field(
+    crn_seed: PRNGKeyArray | None = dataclasses.field(
         default_factory=lambda: jr.PRNGKey(0)
     )
     perturb_measurements: bool | None = None
diff --git a/dynestyx/inference/filters.py b/dynestyx/inference/filters.py
index ae84dc95..e3367f94 100644
--- a/dynestyx/inference/filters.py
+++ b/dynestyx/inference/filters.py
@@ -74,7 +74,7 @@
     FunctionOfTime,
     chain_numpyro_site_registrations,
 )
-from dynestyx.utils import _dist_has_plate_batch_dims
+from dynestyx.utils import _dist_has_plate_batch_dims, _ensure_trailing_event_axis
 
 type SSMType = ContDiscreteNonlinearGaussianSSM | ContDiscreteNonlinearSSM
 
@@ -88,7 +88,7 @@ def _sample_ds(
         name: str,
         dynamics: DynamicalModel,
         *,
-        plate_shapes=(),
+        plate_shapes: tuple[int, ...] = (),
         obs_times: Real[Array, "*obs_time_plate obs_time"] | None = None,
         obs_values: Real[Array, "*obs_value_plate obs_time observation_dim"]
         | Real[Array, "*obs_value_plate obs_time"]
@@ -143,7 +143,7 @@ def _add_log_factors(
         name: str,
         dynamics: DynamicalModel,
         *,
-        plate_shapes=(),
+        plate_shapes: tuple[int, ...] = (),
         obs_times: Real[Array, "*obs_time_plate obs_time"] | None = None,
         obs_values: Real[Array, "*obs_value_plate obs_time observation_dim"]
         | Real[Array, "*obs_value_plate obs_time"]
@@ -161,7 +161,7 @@ def _build_infer_result(
     ) -> ConditionedResult: ...
 
 
-def _default_filter_config(dynamics: DynamicalModel):
+def _default_filter_config(dynamics: DynamicalModel) -> BaseFilterConfig:
     """Return appropriate default filter config when none specified."""
     if dynamics.continuous_time:
         return ContinuousTimeEnKFConfig()
@@ -237,7 +237,7 @@ def _add_log_factors(
         name: str,
         dynamics: DynamicalModel,
         *,
-        plate_shapes=(),
+        plate_shapes: tuple[int, ...] = (),
         obs_times: Real[Array, "*obs_time_plate obs_time"] | None = None,
         obs_values: Real[Array, "*obs_value_plate obs_time observation_dim"]
         | Real[Array, "*obs_value_plate obs_time"]
@@ -275,7 +275,6 @@ def _add_log_factors(
                 obs_values=obs_values,
                 mode="filter",
             )
-
         # Resolve PRNG key: use explicit seed from config, fall back to numpyro
         # context (inside a seeded model), or None (deterministic filters don't need one).
         if config.crn_seed is not None:
@@ -303,6 +302,11 @@ def _add_log_factors(
                 ctrl_values=ctrl_values,
             )
 
+        if not isinstance(config, HMMConfigs):
+            obs_values = _ensure_trailing_event_axis(obs_values)
+            if ctrl_values is not None:
+                ctrl_values = _ensure_trailing_event_axis(ctrl_values)
+
         if dynamics.continuous_time:
             if not isinstance(config, ContinuousTimeConfigs):
                 valid = [c.__name__ for c in ContinuousTimeConfigs]
@@ -439,7 +443,7 @@ def _add_log_factors_batched(
                 )
             output_kind = "continuous"
 
-            def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
+            def _compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
                 return compute_continuous_filter(
                     dyn,
                     cast(ContinuousTimeFilterConfig, config),
@@ -454,7 +458,7 @@ def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
             output_kind = "hmm"
             uses_preprocessed_obs = True
 
-            def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
+            def _compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
                 return compute_hmm_filter(
                     dyn,
                     obs_times=ot,
@@ -468,7 +472,7 @@ def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
             if config.filter_source == "cuthbert":
                 output_kind = "cuthbert"
 
-                def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
+                def _compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
                     return compute_cuthbert_filter(
                         dyn,
                         config,
@@ -482,7 +486,7 @@ def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
             elif config.filter_source == "cd_dynamax":
                 output_kind = "cd_dynamax_discrete"
 
-                def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
+                def _compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
                     return compute_cd_dynamax_discrete_filter(
                         dyn,
                         config,
@@ -499,6 +503,14 @@ def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
                 f"Unsupported filter config for plate: {type(config).__name__}"
             )
 
+        def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
+            # Add scalar event axes after vmap removes plate dimensions.
+            if not isinstance(config, HMMConfigs):
+                ov = _ensure_trailing_event_axis(ov)
+                if cv is not None:
+                    cv = _ensure_trailing_event_axis(cv)
+            return _compute_output(dyn, ot, ov, ovf, om, ct, cv, k)
+
         # Pre-split keys for all plate members (needed for stochastic filters).
         if key is not None:
             # Ensure we use typed PRNG keys so split returns shape (total,)
@@ -669,13 +681,15 @@ def _filter_discrete_time(
     key: PRNGKeyArray | None = None,
     *,
     obs_times: Real[Array, " obs_time"],
-    obs_values: Real[Array, "obs_time observation_dim"] | Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
     ctrl_times: Real[Array, " ctrl_time"] | None = None,
-    ctrl_values: Real[Array, "ctrl_time control_dim"]
-    | Real[Array, " ctrl_time"]
-    | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
     **kwargs,
-) -> tuple[jax.Array | None, object | None, list[numpyro.distributions.Distribution]]:
+) -> tuple[
+    Real[Array, ""] | None,
+    object | None,
+    list[numpyro.distributions.Distribution],
+]:
     """Discrete-time marginal likelihood via cuthbert or cd-dynamax.
 
     Filter type inferred from config class: KFConfig, EKFConfig, UKFConfig
@@ -725,13 +739,15 @@ def _filter_continuous_time(
     key: PRNGKeyArray | None = None,
     *,
     obs_times: Real[Array, " obs_time"],
-    obs_values: Real[Array, "obs_time observation_dim"] | Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
     ctrl_times: Real[Array, " ctrl_time"] | None = None,
-    ctrl_values: Real[Array, "ctrl_time control_dim"]
-    | Real[Array, " ctrl_time"]
-    | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
     **kwargs,
-) -> tuple[jax.Array, object, list[numpyro.distributions.Distribution]]:
+) -> tuple[
+    Real[Array, ""],
+    object,
+    list[numpyro.distributions.Distribution],
+]:
     """Continuous-time marginal likelihood via CD-Dynamax.
 
     Supports: EnKF, DPF, EKF, UKF (inferred from config type).
diff --git a/dynestyx/inference/integrations/cd_dynamax/continuous_filter.py b/dynestyx/inference/integrations/cd_dynamax/continuous_filter.py
index 5756e642..727adf93 100644
--- a/dynestyx/inference/integrations/cd_dynamax/continuous_filter.py
+++ b/dynestyx/inference/integrations/cd_dynamax/continuous_filter.py
@@ -1,5 +1,7 @@
 """Continuous-time filters via CD-Dynamax: KF, EnKF, DPF, EKF, UKF."""
 
+from typing import Any
+
 import jax
 import jax.numpy as jnp
 import numpyro.distributions as dist
@@ -11,6 +13,7 @@
 from cd_dynamax.src.continuous_discrete_linear_gaussian_ssm.models import (
     PosteriorGSSMFiltered,
 )
+from jaxtyping import Array, PRNGKeyArray, Real
 
 from dynestyx.inference.configs.filter import (
     ContinuousTimeDPFConfig,
@@ -39,12 +42,12 @@
 
 def _config_to_cd_dynamax_filter_kwargs(
     config: ContinuousTimeFilterConfig,
-    params,
-    obs_values,
-    obs_times,
-    ctrl_values,
-    key,
-) -> dict:
+    params: Any,
+    obs_values: Real[Array, "obs_time observation_dim"],
+    obs_times: Real[Array, "obs_time one"],
+    ctrl_values: Real[Array, "ctrl_time control_dim"],
+    key: PRNGKeyArray | None,
+) -> dict[str, Any]:
     """Build the filter_kwargs dict passed to cd_dynamax_model.filter()."""
 
     # cd-dynamax uses the legacy PRNG key interface, but newer numpyro uses typed keys.
@@ -111,9 +114,9 @@ def _config_to_cd_dynamax_filter_kwargs(
 
 def _run_linear_kf(
     dynamics: DynamicalModel,
-    obs_times,
-    obs_values,
-    ctrl_values,
+    obs_times: Real[Array, "obs_time one"],
+    obs_values: Real[Array, "obs_time observation_dim"],
+    ctrl_values: Real[Array, "ctrl_time control_dim"],
     filter_config: ContinuousTimeKFConfig,
 ) -> PosteriorGSSMFiltered:
     """Run exact continuous-discrete KF (AffineLinearDrift + constant diffusion + LinearGaussianObservation)."""
@@ -136,13 +139,13 @@ def _run_linear_kf(
 def compute_continuous_filter(
     dynamics: DynamicalModel,
     filter_config: ContinuousTimeFilterConfig,
-    key: jax.Array | None = None,
+    key: PRNGKeyArray | None = None,
     *,
-    obs_times: jax.Array,
-    obs_values: jax.Array,
-    ctrl_times=None,
-    ctrl_values=None,
-):
+    obs_times: Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
+    ctrl_times: Real[Array, " ctrl_time"] | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
+) -> Any:
     """Pure-JAX continuous-time filter computation (no numpyro side-effects)."""
     obs_times_arr = jnp.asarray(obs_times)
     if obs_times_arr.ndim == 1:
@@ -196,14 +199,14 @@ def run_continuous_filter(
     name: str,
     dynamics: DynamicalModel,
     filter_config: ContinuousTimeFilterConfig,
-    key: jax.Array | None = None,
+    key: PRNGKeyArray | None = None,
     *,
-    obs_times: jax.Array,
-    obs_values: jax.Array,
-    ctrl_times=None,
-    ctrl_values=None,
+    obs_times: Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
+    ctrl_times: Real[Array, " ctrl_time"] | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
     **kwargs,
-) -> tuple[jax.Array, object, list[dist.Distribution]]:
+) -> tuple[Real[Array, ""], object, list[dist.Distribution]]:
     """Run continuous-time filter via CD-Dynamax.
 
     Pure computation — no numpyro side-effects. Callers are responsible for
diff --git a/dynestyx/inference/integrations/cd_dynamax/continuous_smoother.py b/dynestyx/inference/integrations/cd_dynamax/continuous_smoother.py
index afd63539..96d5d699 100644
--- a/dynestyx/inference/integrations/cd_dynamax/continuous_smoother.py
+++ b/dynestyx/inference/integrations/cd_dynamax/continuous_smoother.py
@@ -1,5 +1,7 @@
 """Continuous-time smoothers via CD-Dynamax: KF, EKF."""
 
+from typing import Any
+
 import jax
 import jax.numpy as jnp
 import numpyro.distributions as dist
@@ -10,6 +12,7 @@
     cdlgssm_smoother,
     cdnlgssm_smoother,
 )
+from jaxtyping import Array, PRNGKeyArray, Real
 
 from dynestyx.inference.configs.smoother import (
     ContinuousTimeEKFSmootherConfig,
@@ -30,13 +33,13 @@
 def compute_continuous_smoother(
     dynamics: DynamicalModel,
     smoother_config: ContinuousTimeSmootherConfig,
-    key: jax.Array | None = None,
+    key: PRNGKeyArray | None = None,
     *,
-    obs_times: jax.Array,
-    obs_values: jax.Array,
-    ctrl_times=None,
-    ctrl_values=None,
-):
+    obs_times: Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
+    ctrl_times: Real[Array, " ctrl_time"] | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
+) -> Any:
     """Pure-JAX continuous-time smoother computation (no numpyro side-effects)."""
     obs_times_arr = jnp.asarray(obs_times)
     if obs_times_arr.ndim == 1:
@@ -118,14 +121,14 @@ def run_continuous_smoother(
     name: str,
     dynamics: DynamicalModel,
     smoother_config: ContinuousTimeSmootherConfig,
-    key: jax.Array | None = None,
+    key: PRNGKeyArray | None = None,
     *,
-    obs_times: jax.Array,
-    obs_values: jax.Array,
-    ctrl_times=None,
-    ctrl_values=None,
+    obs_times: Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
+    ctrl_times: Real[Array, " ctrl_time"] | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
     **kwargs,
-) -> tuple[jax.Array, object, list[dist.Distribution]]:
+) -> tuple[Real[Array, ""], object, list[dist.Distribution]]:
     """Run continuous-time smoother via CD-Dynamax.
 
     Pure computation — no numpyro side-effects. Callers are responsible for
diff --git a/dynestyx/inference/integrations/cd_dynamax/discrete_filter.py b/dynestyx/inference/integrations/cd_dynamax/discrete_filter.py
index 089c3321..4d7b75be 100644
--- a/dynestyx/inference/integrations/cd_dynamax/discrete_filter.py
+++ b/dynestyx/inference/integrations/cd_dynamax/discrete_filter.py
@@ -1,6 +1,7 @@
 """Discrete-time filters via cd-dynamax (dynamax): KF, EKF, UKF."""
 
-import jax
+from typing import Any, cast
+
 import jax.numpy as jnp
 import numpyro.distributions as dist
 from cd_dynamax.dynamax.linear_gaussian_ssm.inference import (
@@ -14,6 +15,7 @@
     UKFHyperParams,
     unscented_kalman_filter,
 )
+from jaxtyping import Array, Real
 
 from dynestyx.inference.configs.filter import (
     BaseFilterConfig,
@@ -83,7 +85,16 @@ def _lti_to_lgssm_params(dynamics: DynamicalModel):
     )
 
 
-def _prepare_inputs(dynamics, obs_values, obs_times, ctrl_times, ctrl_values):
+def _prepare_inputs(
+    dynamics: DynamicalModel,
+    obs_values: Real[Array, "obs_time observation_dim"],
+    obs_times: Real[Array, " obs_time"],
+    ctrl_times: Real[Array, " ctrl_time"] | None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None,
+) -> tuple[
+    Real[Array, "obs_time observation_dim"],
+    Real[Array, "obs_time control_dim"],
+]:
     """Prepare emissions and inputs arrays for cd-dynamax discrete filters."""
     emissions = obs_values
     t1 = emissions.shape[0]
@@ -91,7 +102,8 @@ def _prepare_inputs(dynamics, obs_values, obs_times, ctrl_times, ctrl_values):
     if ctrl_values is None:
         inputs = jnp.zeros((t1, control_dim))
     elif ctrl_values.shape[0] > t1:
-        inds = jnp.searchsorted(ctrl_times, obs_times, side="left")
+        aligned_ctrl_times = cast(Real[Array, " ctrl_time"], ctrl_times)
+        inds = jnp.searchsorted(aligned_ctrl_times, obs_times, side="left")
         inputs = ctrl_values[inds]
     else:
         inputs = ctrl_values
@@ -102,11 +114,11 @@ def compute_cd_dynamax_discrete_filter(
     dynamics: DynamicalModel,
     filter_config: BaseFilterConfig,
     *,
-    obs_times: jax.Array,
-    obs_values: jax.Array,
-    ctrl_times=None,
-    ctrl_values=None,
-):
+    obs_times: Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
+    ctrl_times: Real[Array, " ctrl_time"] | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
+) -> Any:
     """Pure-JAX cd-dynamax discrete filter computation (no numpyro side-effects)."""
     emissions, inputs = _prepare_inputs(
         dynamics, obs_values, obs_times, ctrl_times, ctrl_values
@@ -141,12 +153,12 @@ def run_discrete_filter(
     dynamics: DynamicalModel,
     filter_config: BaseFilterConfig,
     *,
-    obs_times: jax.Array,
-    obs_values: jax.Array,
-    ctrl_times=None,
-    ctrl_values=None,
+    obs_times: Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
+    ctrl_times: Real[Array, " ctrl_time"] | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
     **kwargs,
-) -> tuple[jax.Array, object, list[dist.Distribution]]:
+) -> tuple[Real[Array, ""], object, list[dist.Distribution]]:
     """Run discrete-time filter via cd-dynamax (KF, EKF, UKF).
 
     Pure computation — no numpyro side-effects. Callers are responsible for
diff --git a/dynestyx/inference/integrations/cd_dynamax/discrete_smoother.py b/dynestyx/inference/integrations/cd_dynamax/discrete_smoother.py
index afa53289..84832091 100644
--- a/dynestyx/inference/integrations/cd_dynamax/discrete_smoother.py
+++ b/dynestyx/inference/integrations/cd_dynamax/discrete_smoother.py
@@ -1,6 +1,7 @@
 """Discrete-time smoothers via cd-dynamax (dynamax): KF, EKF, UKF."""
 
-import jax
+from typing import Any
+
 import numpyro.distributions as dist
 from cd_dynamax.dynamax.linear_gaussian_ssm.inference import lgssm_smoother
 from cd_dynamax.dynamax.nonlinear_gaussian_ssm.inference_ekf import (
@@ -10,6 +11,7 @@
     UKFHyperParams,
     unscented_kalman_smoother,
 )
+from jaxtyping import Array, Real
 
 from dynestyx.inference.configs.filter import (
     EKFConfig,
@@ -32,11 +34,11 @@ def compute_cd_dynamax_discrete_smoother(
     dynamics: DynamicalModel,
     filter_config: BaseSmootherConfig,
     *,
-    obs_times: jax.Array,
-    obs_values: jax.Array,
-    ctrl_times=None,
-    ctrl_values=None,
-):
+    obs_times: Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
+    ctrl_times: Real[Array, " ctrl_time"] | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
+) -> Any:
     """Pure-JAX cd-dynamax discrete smoother computation (no numpyro side-effects)."""
     emissions, inputs = _prepare_inputs(
         dynamics, obs_values, obs_times, ctrl_times, ctrl_values
@@ -69,12 +71,12 @@ def run_discrete_smoother(
     dynamics: DynamicalModel,
     filter_config: BaseSmootherConfig,
     *,
-    obs_times: jax.Array,
-    obs_values: jax.Array,
-    ctrl_times=None,
-    ctrl_values=None,
+    obs_times: Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
+    ctrl_times: Real[Array, " ctrl_time"] | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
     **kwargs,
-) -> tuple[jax.Array, object, list[dist.Distribution]]:
+) -> tuple[Real[Array, ""], object, list[dist.Distribution]]:
     """Run discrete-time smoother via cd-dynamax (KF, EKF, UKF).
 
     Pure computation — no numpyro side-effects. Callers are responsible for
diff --git a/dynestyx/inference/integrations/cuthbert/discrete_filter.py b/dynestyx/inference/integrations/cuthbert/discrete_filter.py
index aef1c4b2..39ee780c 100644
--- a/dynestyx/inference/integrations/cuthbert/discrete_filter.py
+++ b/dynestyx/inference/integrations/cuthbert/discrete_filter.py
@@ -1,5 +1,5 @@
 import warnings
-from typing import NamedTuple, cast
+from typing import Any, NamedTuple, cast
 
 import jax
 import jax.numpy as jnp
@@ -14,6 +14,7 @@
     stop_gradient_decorator,
     systematic,
 )
+from jaxtyping import Array, Bool, Float, PRNGKeyArray, Real
 
 from dynestyx.inference.configs.filter import (
     BaseFilterConfig,
@@ -37,17 +38,19 @@
 
 
 class CuthbertInputs(NamedTuple):
-    """Model inputs pytree for cuthbert; leading time dim must be T+1."""
+    """Model-input pytree before or after cuthbert slices its leading time axis."""
 
-    y: jax.Array  # (T+1, emission_dim)
-    u: jax.Array  # (T+1, control_dim) or (T+1, 0)
-    u_prev: jax.Array  # (T+1, control_dim) or (T+1, 0)
-    time: jax.Array  # (T+1,)
-    time_prev: jax.Array  # (T+1,)
-    is_first_step: jax.Array  # (T+1,) bool — True only at index 1
+    y: Real[Array, "cuthbert_time observation_dim"] | Real[Array, " observation_dim"]
+    u: Real[Array, "cuthbert_time control_dim"] | Real[Array, " control_dim"]
+    u_prev: Real[Array, "cuthbert_time control_dim"] | Real[Array, " control_dim"]
+    time: Real[Array, " cuthbert_time"] | Real[Array, ""]
+    time_prev: Real[Array, " cuthbert_time"] | Real[Array, ""]
+    is_first_step: Bool[Array, " cuthbert_time"] | Bool[Array, ""]
 
 
-def _extract_gaussian_chol(d: dist.Distribution, obs_dim: int) -> jax.Array:
+def _extract_gaussian_chol(
+    d: dist.Distribution, obs_dim: int
+) -> Float[Array, "observation_dim observation_dim"]:
     """Extract a Cholesky factor of the covariance from a Gaussian distribution."""
     if isinstance(d, dist.MultivariateNormal):
         return jnp.asarray(d.scale_tril)
@@ -68,7 +71,9 @@ def _extract_gaussian_chol(d: dist.Distribution, obs_dim: int) -> jax.Array:
 
 
 def _check_state_independent_noise(
-    chol_R_at_x0: jax.Array, probe_dist_at_x1: dist.Distribution, obs_dim: int
+    chol_R_at_x0: Float[Array, "observation_dim observation_dim"],
+    probe_dist_at_x1: dist.Distribution,
+    obs_dim: int,
 ) -> None:
     """Raise if the observation noise covariance varies with state."""
     chol_R_at_x1 = _extract_gaussian_chol(probe_dist_at_x1, obs_dim)
@@ -154,14 +159,14 @@ def _drop_if_time_leaf(leaf):
 def compute_cuthbert_filter(
     dynamics: DynamicalModel,
     filter_config: BaseFilterConfig,
-    key: jax.Array | None = None,
+    key: PRNGKeyArray | None = None,
     *,
-    obs_times: jax.Array,
-    obs_values: jax.Array,
-    ctrl_times=None,
-    ctrl_values=None,
+    obs_times: Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
+    ctrl_times: Real[Array, " ctrl_time"] | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
     align_to_observations: bool = True,
-):
+) -> tuple[Real[Array, ""], Any]:
     """Pure-JAX cuthbert filter computation (no numpyro side-effects).
 
     Returns:
@@ -250,14 +255,14 @@ def run_discrete_filter(
     name: str,
     dynamics: DynamicalModel,
     filter_config: BaseFilterConfig,
-    key: jax.Array | None = None,
+    key: PRNGKeyArray | None = None,
     *,
-    obs_times: jax.Array,
-    obs_values: jax.Array,
-    ctrl_times=None,
-    ctrl_values=None,
+    obs_times: Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
+    ctrl_times: Real[Array, " ctrl_time"] | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
     **kwargs,
-) -> tuple[jax.Array | None, object | None, list[dist.Distribution]]:
+) -> tuple[Real[Array, ""] | None, object | None, list[dist.Distribution]]:
     """Run discrete-time filter via cuthbert (Kalman, Taylor KF, particle filter).
 
     Pure computation — no numpyro side-effects. Callers are responsible for
diff --git a/dynestyx/inference/integrations/cuthbert/discrete_smoother.py b/dynestyx/inference/integrations/cuthbert/discrete_smoother.py
index 75274656..62bcfd73 100644
--- a/dynestyx/inference/integrations/cuthbert/discrete_smoother.py
+++ b/dynestyx/inference/integrations/cuthbert/discrete_smoother.py
@@ -2,9 +2,8 @@
 
 from collections.abc import Callable
 from functools import partial
-from typing import cast
+from typing import Any, cast
 
-import jax
 import jax.numpy as jnp
 import numpyro.distributions as dist
 from cuthbert import smoother as cuthbert_smoother
@@ -12,6 +11,7 @@
 from cuthbert.smc import backward_sampler
 from cuthbertlib.resampling import multinomial, stop_gradient_decorator, systematic
 from cuthbertlib.smc.smoothing import exact_sampling, mcmc, tracing
+from jaxtyping import Array, PRNGKeyArray, Real
 
 from dynestyx.inference.configs.smoother import (
     EKFSmootherConfig,
@@ -66,7 +66,13 @@ def _kalman_get_dynamics_params(dynamics: DynamicalModel):
 def _taylor_get_dynamics_log_density(dynamics: DynamicalModel):
     transition = cast(
         Callable[
-            [jax.Array, jax.Array | None, jax.Array, jax.Array], dist.Distribution
+            [
+                Real[Array, " state_dim"],
+                Real[Array, " control_dim"] | Real[Array, ""] | None,
+                Real[Array, ""],
+                Real[Array, ""],
+            ],
+            dist.Distribution,
         ],
         dynamics.state_evolution,
     )
@@ -150,13 +156,13 @@ def _pf_backward_sampling_fn(config: PFSmootherConfig):
 def compute_cuthbert_smoother(
     dynamics: DynamicalModel,
     smoother_config: CuthbertSmootherConfig,
-    key: jax.Array | None = None,
+    key: PRNGKeyArray | None = None,
     *,
-    obs_times: jax.Array,
-    obs_values: jax.Array,
-    ctrl_times=None,
-    ctrl_values=None,
-):
+    obs_times: Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
+    ctrl_times: Real[Array, " ctrl_time"] | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
+) -> tuple[Real[Array, ""], Any]:
     """Pure-JAX cuthbert smoother computation (no numpyro side-effects)."""
     obs_len = int(obs_values.shape[0])
     marginal_loglik, filtered_states = compute_cuthbert_filter(
@@ -226,14 +232,14 @@ def run_discrete_smoother(
     name: str,
     dynamics: DynamicalModel,
     smoother_config: CuthbertSmootherConfig,
-    key: jax.Array | None = None,
+    key: PRNGKeyArray | None = None,
     *,
-    obs_times: jax.Array,
-    obs_values: jax.Array,
-    ctrl_times=None,
-    ctrl_values=None,
+    obs_times: Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
+    ctrl_times: Real[Array, " ctrl_time"] | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
     **kwargs,
-) -> tuple[jax.Array | None, object | None, list[dist.Distribution]]:
+) -> tuple[Real[Array, ""] | None, object | None, list[dist.Distribution]]:
     """Run discrete-time smoother via cuthbert.
 
     Returns:
diff --git a/dynestyx/inference/smoothers.py b/dynestyx/inference/smoothers.py
index e30ff6e6..c1db69e8 100644
--- a/dynestyx/inference/smoothers.py
+++ b/dynestyx/inference/smoothers.py
@@ -65,7 +65,7 @@
     FunctionOfTime,
     chain_numpyro_site_registrations,
 )
-from dynestyx.utils import _dist_has_plate_batch_dims
+from dynestyx.utils import _dist_has_plate_batch_dims, _ensure_trailing_event_axis
 
 DiscreteSmootherConfig = (
     KFSmootherConfig | EKFSmootherConfig | UKFSmootherConfig | PFSmootherConfig
@@ -98,7 +98,7 @@ def _sample_ds(
         name: str,
         dynamics: DynamicalModel,
         *,
-        plate_shapes=(),
+        plate_shapes: tuple[int, ...] = (),
         obs_times: Real[Array, "*obs_time_plate obs_time"] | None = None,
         obs_values: Real[Array, "*obs_value_plate obs_time observation_dim"]
         | Real[Array, "*obs_value_plate obs_time"]
@@ -179,7 +179,7 @@ def _add_log_factors(
         name: str,
         dynamics: DynamicalModel,
         *,
-        plate_shapes=(),
+        plate_shapes: tuple[int, ...] = (),
         obs_times: Real[Array, "*obs_time_plate obs_time"] | None = None,
         obs_values: Real[Array, "*obs_value_plate obs_time observation_dim"]
         | Real[Array, "*obs_value_plate obs_time"]
@@ -202,7 +202,7 @@ class Smoother(BaseSmootherLogFactorAdder):
     r"""Performs Bayesian smoothing to compute the smoothing distribution p(x_t | y_{1:T})."""
 
     smoother_config: SmootherAnyConfig | None = None
-    marginal_loglik: jax.Array | None = dataclasses.field(
+    marginal_loglik: Real[Array, "*plate"] | None = dataclasses.field(
         default=None, repr=False, init=False
     )
     smoothed_states: object = dataclasses.field(default=None, repr=False, init=False)
@@ -215,7 +215,7 @@ def _add_log_factors(
         name: str,
         dynamics: DynamicalModel,
         *,
-        plate_shapes=(),
+        plate_shapes: tuple[int, ...] = (),
         obs_times: Real[Array, "*obs_time_plate obs_time"] | None = None,
         obs_values: Real[Array, "*obs_value_plate obs_time observation_dim"]
         | Real[Array, "*obs_value_plate obs_time"]
@@ -249,7 +249,6 @@ def _add_log_factors(
             obs_values=obs_values,
             mode="smoother",
         )
-
         # Resolve PRNG key: use explicit seed from config, fall back to numpyro
         # context (inside a seeded model), or None (deterministic smoothers don't need one).
         typed_config = config
@@ -275,6 +274,10 @@ def _add_log_factors(
                 ctrl_values=ctrl_values,
             )
 
+        obs_values = _ensure_trailing_event_axis(obs_values)
+        if ctrl_values is not None:
+            ctrl_values = _ensure_trailing_event_axis(ctrl_values)
+
         if dynamics.continuous_time:
             if not isinstance(typed_config, ContinuousTimeSmootherConfigs):
                 valid = _valid_smoother_config_names(continuous_time=True)
@@ -376,7 +379,7 @@ def _add_log_factors_batched(
             continuous_config = cast(ContinuousSmootherConfig, config)
             output_kind = "continuous"
 
-            def compute_output(dyn, ot, ov, ct, cv, k):
+            def _compute_output(dyn, ot, ov, ct, cv, k):
                 return compute_continuous_smoother(
                     dyn,
                     continuous_config,
@@ -398,7 +401,7 @@ def compute_output(dyn, ot, ov, ct, cv, k):
                 cuthbert_config = cast(CuthbertSmootherConfig, discrete_config)
                 output_kind = "cuthbert"
 
-                def compute_output(dyn, ot, ov, ct, cv, k):
+                def _compute_output(dyn, ot, ov, ct, cv, k):
                     return compute_cuthbert_smoother(
                         dyn,
                         cuthbert_config,
@@ -412,7 +415,7 @@ def compute_output(dyn, ot, ov, ct, cv, k):
             elif discrete_config.filter_source == "cd_dynamax":
                 output_kind = "cd_dynamax_discrete"
 
-                def compute_output(dyn, ot, ov, ct, cv, k):
+                def _compute_output(dyn, ot, ov, ct, cv, k):
                     return compute_cd_dynamax_discrete_smoother(
                         dyn,
                         discrete_config,
@@ -431,6 +434,13 @@ def compute_output(dyn, ot, ov, ct, cv, k):
                 f"Unsupported smoother config for plate: {type(config).__name__}"
             )
 
+        def compute_output(dyn, ot, ov, ct, cv, k):
+            # Add scalar event axes after vmap removes plate dimensions.
+            ov = _ensure_trailing_event_axis(ov)
+            if cv is not None:
+                cv = _ensure_trailing_event_axis(cv)
+            return _compute_output(dyn, ot, ov, ct, cv, k)
+
         if key is not None:
             if not jnp.issubdtype(key.dtype, jax.dtypes.prng_key):
                 key = jax.random.wrap_key_data(key)
@@ -559,13 +569,15 @@ def _smooth_discrete_time(
     key: PRNGKeyArray | None = None,
     *,
     obs_times: Real[Array, " obs_time"],
-    obs_values: Real[Array, "obs_time observation_dim"] | Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
     ctrl_times: Real[Array, " ctrl_time"] | None = None,
-    ctrl_values: Real[Array, "ctrl_time control_dim"]
-    | Real[Array, " ctrl_time"]
-    | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
     **kwargs,
-) -> tuple[jax.Array | None, object | None, list[numpyro.distributions.Distribution]]:
+) -> tuple[
+    Real[Array, ""] | None,
+    object | None,
+    list[numpyro.distributions.Distribution],
+]:
     """Discrete-time marginal likelihood via cuthbert or cd-dynamax smoothers."""
 
     if isinstance(smoother_config, UKFSmootherConfig) and (
@@ -626,13 +638,15 @@ def _smooth_continuous_time(
     key: PRNGKeyArray | None = None,
     *,
     obs_times: Real[Array, " obs_time"],
-    obs_values: Real[Array, "obs_time observation_dim"] | Real[Array, " obs_time"],
+    obs_values: Real[Array, "obs_time observation_dim"],
     ctrl_times: Real[Array, " ctrl_time"] | None = None,
-    ctrl_values: Real[Array, "ctrl_time control_dim"]
-    | Real[Array, " ctrl_time"]
-    | None = None,
+    ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
     **kwargs,
-) -> tuple[jax.Array, object, list[numpyro.distributions.Distribution]]:
+) -> tuple[
+    Real[Array, ""],
+    object,
+    list[numpyro.distributions.Distribution],
+]:
     """Continuous-time marginal likelihood via CD-Dynamax smoothers."""
     if smoother_config.filter_source != "cd_dynamax":
         raise ValueError(
diff --git a/dynestyx/utils.py b/dynestyx/utils.py
index 97a8ee4b..75d70277 100644
--- a/dynestyx/utils.py
+++ b/dynestyx/utils.py
@@ -44,6 +44,15 @@ def flatten_draws(arr: Shaped[Array, "..."]) -> Shaped[Array, "..."]:
 _CONTROL_EXTEND_EPSILON = 1e-5
 
 
+def _ensure_trailing_event_axis(
+    values: Real[Array, "..."],
+) -> Real[Array, "..."]:
+    """Lift a scalar time series from ``(time,)`` to ``(time, 1)``."""
+    if values.ndim == 1:
+        return values[..., None]
+    return values
+
+
 def _raise_now_or_error_if(
     anchor: Array,
     predicate,