Idiomatic Jax for linear filter, though no speed difference. Added a note.

Author: hmgaudecker

docs/explanations/linear_predict.md

@@ -50,6 +50,19 @@ For each latent factor $i$:
 - **Constant factor**: row $i$ of $F$ is the unit vector $e_i$ (identity row) and
   $c_i = 0$, so the factor value is simply carried forward.
 
+The implementation uses a stack-then-mask approach: all coefficient arrays are stacked
+into a single matrix (with zero-padded rows for constant factors), an identity matrix
+provides the constant-factor rows, and `jnp.where` selects between them using a boolean
+mask. This avoids per-element `.at[i].set()` calls and conditional branching, producing
+a cleaner trace for JAX's compiler.
+
+Three construction strategies were benchmarked (loop with conditional `.at[i].set()`,
+stack-then-mask with `jnp.where`, and index-scatter with pre-separated sub-matrices).
+All three produced identical XLA graphs and showed no meaningful runtime difference
+(~6.3--6.7 ms per call on CPU, 4-factor model, 5000 observations), confirming that the
+construction is fully resolved at trace time. The stack-then-mask variant was kept for
+its cleaner, more idiomatic JAX style.
+
 ## Mean prediction
 
 The mean prediction incorporates anchoring, which rescales factors to a common metric

src/skillmodels/kalman_filters.py

@@ -284,23 +284,9 @@ def linear_kalman_predict(
     """
     n_latent = len(latent_factors)
 
-    # Build F (n_latent x n_all) and c (n_latent,) from trans_coeffs.
-    # linear factor i: F[i] = trans_coeffs[factor_i][:-1], c[i] = last element
-    # constant factor i: F[i] = e_i (unit vector), c[i] = 0
-    f_rows = []
-    c_vals = []
-    for i, factor in enumerate(latent_factors):
-        if i in constant_factor_indices:
-            row = jnp.zeros(n_all_factors).at[i].set(1.0)
-            f_rows.append(row)
-            c_vals.append(0.0)
-        else:
-            coeffs = trans_coeffs[factor]
-            f_rows.append(coeffs[:-1])
-            c_vals.append(coeffs[-1])
-
-    f_mat = jnp.stack(f_rows)  # (n_latent, n_all)
-    c_vec = jnp.array(c_vals)  # (n_latent,)
+    f_mat, c_vec = _build_f_and_c(
+        latent_factors, constant_factor_indices, n_all_factors, trans_coeffs
+    )
 
     s_in = anchoring_scaling_factors[0][:n_latent]  # (n_latent,) for input period
     s_out = anchoring_scaling_factors[1][:n_latent]  # (n_latent,) for output period
@@ -338,6 +324,53 @@ def linear_kalman_predict(
     return predicted_states, predicted_covs
 
 
+def _build_f_and_c(
+    latent_factors: tuple[str, ...],
+    constant_factor_indices: frozenset[int],
+    n_all_factors: int,
+    trans_coeffs: dict[str, Array],
+) -> tuple[Array, Array]:
+    """Build F matrix and c vector from transition coefficients.
+
+    Stack all coefficient arrays, build identity rows for constant factors,
+    and select via a boolean mask.
+
+    Args:
+        latent_factors: Tuple of latent factor names.
+        constant_factor_indices: Indices of factors with `constant` transition.
+        n_all_factors: Total number of factors (latent + observed).
+        trans_coeffs: Dict mapping factor name to 1d coefficient array.
+
+    Returns:
+        f_mat: Array of shape (n_latent, n_all_factors).
+        c_vec: Array of shape (n_latent,).
+
+    """
+    n_latent = len(latent_factors)
+    identity = jnp.eye(n_latent, n_all_factors)
+
+    # Will be of shape (n_latent, n_all+1)
+    all_coeffs = jnp.stack(
+        [
+            trans_coeffs[f]
+            if i not in constant_factor_indices
+            else jnp.zeros(n_all_factors + 1)
+            for i, f in enumerate(latent_factors)
+        ]
+    )
+
+    f_from_coeffs = all_coeffs[:, :-1]
+    c_from_coeffs = all_coeffs[:, -1]
+
+    is_constant = jnp.array([i in constant_factor_indices for i in range(n_latent)])
+    mask = is_constant[:, None]  # (n_latent, 1)
+
+    f_mat = jnp.where(mask, identity, f_from_coeffs)
+    c_vec = jnp.where(is_constant, 0.0, c_from_coeffs)
+
+    return f_mat, c_vec
+
+
 def _calculate_sigma_points(
     states: Array,
     upper_chols: Array,