Fix compile slowdown in filtering

dynamax/generalized_gaussian_ssm/inference.py

@@ -320,7 +320,7 @@ def conditional_moments_gaussian_filter(model_params: ParamsGGSSM,
     # Emission distribution
     emission_dist = model_params.emission_dist
 
-    def _step(carry, t, predict=True):
+    def _step(carry, t):
         """One step of the CMGF"""
         ll, pred_mean, pred_cov = carry
 
@@ -334,17 +334,15 @@ def _step(carry, t, predict=True):
         ll += log_likelihood
 
         # Predict the next state
-        if predict:
-            Q = _get_params(model_params.dynamics_covariance, 2, t + 1)
-            pred_mean, pred_cov, _ = _predict(filtered_mean, filtered_cov, f, Q, inputs[t + 1], g_ev, g_cov)
+        t_next = jnp.minimum(t + 1, num_timesteps - 1)
+        Q = _get_params(model_params.dynamics_covariance, 2, t_next)
+        pred_mean, pred_cov, _ = _predict(filtered_mean, filtered_cov, f, Q, inputs[t_next], g_ev, g_cov)
 
         return (ll, pred_mean, pred_cov), (filtered_mean, filtered_cov)
 
     # Run the general linearization filter
     carry = (0.0, model_params.initial_mean, model_params.initial_covariance)
-    carry, outputs = lax.scan(_step, carry, jnp.arange(num_timesteps - 1))
-    (ll, _, _), final_outputs = _step(carry, num_timesteps - 1, predict=False)
-    filtered_means, filtered_covs = [jnp.concatenate((a, b[None])) for a, b in zip(outputs, final_outputs)]
+    (ll, _, _), (filtered_means, filtered_covs) = lax.scan(_step, carry, jnp.arange(num_timesteps))
     return PosteriorGSSMFiltered(marginal_loglik=ll, 
                                  filtered_means=filtered_means, 
                                  filtered_covariances=filtered_covs)

dynamax/linear_gaussian_ssm/inference.py

@@ -496,7 +496,7 @@ def _log_likelihood(pred_mean, pred_cov, H, D, d, R, u, y):
             return MVNLowRank(m, R, L).log_prob(y)
 
 
-    def _step(carry, t, predict=True):
+    def _step(carry, t):
         """Run one step of the Kalman filter."""
         ll, pred_mean, pred_cov = carry
 
@@ -512,18 +512,15 @@ def _step(carry, t, predict=True):
         filtered_mean, filtered_cov = _condition_on(pred_mean, pred_cov, H, D, d, R, u, y)
 
         # Predict the next state
-        if predict:
-            F, B, b, Q = _get_params(params, num_timesteps, t + 1)[:4]
-            pred_mean, pred_cov = _predict(filtered_mean, filtered_cov, F, B, b, Q, inputs[t + 1])
+        t_next = jnp.minimum(t + 1, num_timesteps - 1)
+        F, B, b, Q = _get_params(params, num_timesteps, t_next)[:4]
+        pred_mean, pred_cov = _predict(filtered_mean, filtered_cov, F, B, b, Q, inputs[t_next])
 
         return (ll, pred_mean, pred_cov), (filtered_mean, filtered_cov)
 
     # Run the Kalman filter
     carry = (0.0, params.initial.mean, params.initial.cov)
-    carry, (filtered_means, filtered_covs) = lax.scan(_step, carry, jnp.arange(num_timesteps - 1))
-    (ll, _, _), (filtered_mean, filtered_cov) = _step(carry, num_timesteps - 1, predict=False)
-    filtered_means = jnp.concatenate((filtered_means, filtered_mean[None]))
-    filtered_covs = jnp.concatenate((filtered_covs, filtered_cov[None]))
+    (ll, _, _), (filtered_means, filtered_covs) = lax.scan(_step, carry, jnp.arange(num_timesteps))
     return PosteriorGSSMFiltered(marginal_loglik=ll, filtered_means=filtered_means, filtered_covariances=filtered_covs)
 
 

dynamax/linear_gaussian_ssm/info_inference.py

@@ -192,7 +192,7 @@ def lgssm_info_filter(
         if param is not None and not callable(param) and param.ndim == ndim + 1:
             assert param.shape[0] == num_timesteps
     inputs = jnp.zeros((num_timesteps, 0)) if inputs is None else inputs
-    def _filter_step(carry, t, predict=True):
+    def _filter_step(carry, t):
         """Run a single step of the Kalman filter."""
         ll, pred_eta, pred_prec = carry
 
@@ -212,22 +212,19 @@ def _filter_step(carry, t, predict=True):
         filtered_eta, filtered_prec = _info_condition_on(pred_eta, pred_prec, H, R_prec, D, u, d, y)
 
         # Predict the next state
-        if predict:
-            F = _get_params(params.dynamics_weights, 2, t + 1)
-            Q_prec = _get_params(params.dynamics_precision, 2, t + 1)
-            B = _get_params(params.dynamics_input_weights, 2, t + 1)
-            b = _get_params(params.dynamics_bias, 1, t + 1)
-            pred_eta, pred_prec = _info_predict(filtered_eta, filtered_prec, F, Q_prec, B, inputs[t + 1], b)
+        t_next = jnp.minimum(t + 1, num_timesteps - 1)
+        F = _get_params(params.dynamics_weights, 2, t_next)
+        Q_prec = _get_params(params.dynamics_precision, 2, t_next)
+        B = _get_params(params.dynamics_input_weights, 2, t_next)
+        b = _get_params(params.dynamics_bias, 1, t_next)
+        pred_eta, pred_prec = _info_predict(filtered_eta, filtered_prec, F, Q_prec, B, inputs[t_next], b)
 
         return (ll, pred_eta, pred_prec), (filtered_eta, filtered_prec)
 
     # Run the Kalman filter
     initial_eta = params.initial_precision @ params.initial_mean
     carry = (0.0, initial_eta, params.initial_precision)
-    carry, (filtered_etas, filtered_precisions) = lax.scan(_filter_step, carry, jnp.arange(num_timesteps - 1))
-    (ll, _, _), (filtered_eta, filtered_prec) = _filter_step(carry, num_timesteps - 1, predict=False)
-    filtered_etas = jnp.concatenate((filtered_etas, filtered_eta[None]))
-    filtered_precisions = jnp.concatenate((filtered_precisions, filtered_prec[None]))
+    (ll, _, _), (filtered_etas, filtered_precisions) = lax.scan(_filter_step, carry, jnp.arange(num_timesteps))
     return PosteriorGSSMInfoFiltered(marginal_loglik=ll, filtered_etas=filtered_etas, filtered_precisions=filtered_precisions)
 
 

dynamax/nonlinear_gaussian_ssm/inference_ekf.py

@@ -121,7 +121,7 @@ def extended_kalman_filter(params: ParamsNLGSSM,
     f, h, F, H = (_process_fn(fn, inputs) for fn in (f, h, F, H))
     inputs = _process_input(inputs, num_timesteps)
 
-    def _step(carry, t, predict=True):
+    def _step(carry, t):
         """Iteratively update the state estimate and log likelihood."""
         ll, pred_mean, pred_cov = carry
 
@@ -138,10 +138,9 @@ def _step(carry, t, predict=True):
         filtered_mean, filtered_cov = _condition_on(pred_mean, pred_cov, h, H, R, u, y, num_iter)
 
         # Predict the next state
-        if predict:
-            Q = _get_params(params.dynamics_covariance, 2, t + 1)
-            u_next = inputs[t + 1] if not predict_final else u
-            pred_mean, pred_cov = _predict(filtered_mean, filtered_cov, f, F, Q, u_next)
+        t_next = jnp.minimum(t + 1, num_timesteps - 1)
+        Q = _get_params(params.dynamics_covariance, 2, t_next)
+        pred_mean, pred_cov = _predict(filtered_mean, filtered_cov, f, F, Q, inputs[t_next])
 
         # Build carry and output states
         carry = (ll, pred_mean, pred_cov)
@@ -158,9 +157,7 @@ def _step(carry, t, predict=True):
 
     # Run the extended Kalman filter
     carry = (0.0, params.initial_mean, params.initial_covariance)
-    carry, outputs = lax.scan(_step, carry, jnp.arange(num_timesteps - 1))
-    (ll, _, _), final_outputs = _step(carry, num_timesteps - 1, predict=predict_final)
-    outputs = {key: jnp.concatenate((val, final_outputs[key][None])) for key, val in outputs.items()}
+    (ll, _, _), outputs = lax.scan(_step, carry, jnp.arange(num_timesteps))
     if not predict_final:
         outputs = {key: val[:-1] if key.startswith("predicted") else val for key, val in outputs.items()}
     outputs = {"marginal_loglik": ll, **outputs}

dynamax/nonlinear_gaussian_ssm/inference_ukf.py

@@ -199,7 +199,7 @@ def unscented_kalman_filter(params: ParamsNLGSSM,
     f, h = (_process_fn(fn, inputs) for fn in (f, h))
     inputs = _process_input(inputs, num_timesteps)
 
-    def _step(carry, t, predict=True):
+    def _step(carry, t):
         """One step of the UKF"""
         ll, pred_mean, pred_cov = carry
 
@@ -217,10 +217,9 @@ def _step(carry, t, predict=True):
         ll += log_likelihood
 
         # Predict the next state
-        if predict:
-            Q = _get_params(params.dynamics_covariance, 2, t + 1)
-            u_next = inputs[t + 1] if not predict_final else u
-            pred_mean, pred_cov, _ = _predict(filtered_mean, filtered_cov, f, Q, lamb, w_mean, w_cov, u_next)
+        t_next = jnp.minimum(t + 1, num_timesteps - 1)
+        Q = _get_params(params.dynamics_covariance, 2, t_next)
+        pred_mean, pred_cov, _ = _predict(filtered_mean, filtered_cov, f, Q, lamb, w_mean, w_cov, inputs[t_next])
 
         # Build carry and output states
         carry = (ll, pred_mean, pred_cov)
@@ -237,9 +236,7 @@ def _step(carry, t, predict=True):
 
     # Run the Unscented Kalman Filter
     carry = (0.0, params.initial_mean, params.initial_covariance)
-    carry, outputs = lax.scan(_step, carry, jnp.arange(num_timesteps - 1))
-    (ll, _, _), final_outputs = _step(carry, num_timesteps - 1, predict=predict_final)
-    outputs = {key: jnp.concatenate((val, final_outputs[key][None])) for key, val in outputs.items()}
+    (ll, _, _), outputs = lax.scan(_step, carry, jnp.arange(num_timesteps))
     if not predict_final:
         outputs = {key: val[:-1] if key.startswith("predicted") else val for key, val in outputs.items()}
     outputs = {"marginal_loglik": ll, **outputs}