fix: Enzyme syrk adjoint workaround for Kalman filter

[jlperla]

Summary

• Adds mul_aat!! and muladd_aat!! utility functions that compute Y = A*A' and Y += A*A' without triggering the BLAS syrk path
• Updates alloc_kalman_cache, zero_kalman_cache!!, and kalman! to use the new workaround functions
• Workaround for Enzyme.jl#2355: Enzyme's syrk adjoint produces wrong gradients for rectangular matrices

Test plan

• CI builds the notebook successfully
• All regression test values unchanged (workaround computes same result, different BLAS path)
• Enzyme AD sections execute without errors

🤖 Generated with Claude Code

[github-actions]

🚀 Deployed on https://69bf705e1e6683aefc76230f--vibrant-allen-f39ec9.netlify.app

[jlperla]

@claude review for correctness

[claude]

Claude finished @jlperla's task in 1m 31s —— View job

---

Code Review

• Gather context and read changed files
• Review correctness of mul_aat!! and muladd_aat!! implementations
• Review integration into Kalman filter functions
• Check mathematical correctness
• Post review findings

---

Overall: correct. The workaround is sound and the dimensions check out. A few notes below.

mul_aat!! / muladd_aat!! implementations

The core workaround is correct: passing a lazy transpose(A) to mul! can dispatch to BLAS syrk (exploiting symmetry of the result), while passing a pre-allocated materialized transpose avoids that path and uses gemm instead. The ismutable(Y) dispatch correctly handles the mutable (Matrix) vs immutable (SMatrix) case.

Minor: 1 vs 1.0 inconsistency between doc-table and code.
The table in the docs says mul!(Y,A,A_t,1,1) but the code uses mul!(Y, A, A_t, 1.0, 1.0). For Float64 this is irrelevant. If Float32 arrays were ever used, passing Float64 literals 1.0 could cause a promotion that misses a specialized BLAS path — but this is unlikely to matter in practice here. Fix this →

Buffer dimensions

• C is N × K (state dim × process noise dim), so C' is K × N. Buffer C_t = alloc_like(Sigma_0, K, N) ✓
• H is M × L (obs dim × obs noise dim), so H' is L × M. Buffer H_t = alloc_like(Sigma_0, L, M) ✓

Integration into alloc_kalman_cache

K = size(model.C, 2) and L = size(model.H, 2) correctly extract the noise dimensions, and alloc_like propagates the element type from Sigma_0. Correct.

zero_kalman_cache!!

fill_zero!!(cache.C_t)
fill_zero!!(cache.H_t)

The return values are discarded (unlike cache.mu_pred[t] = fill_zero!!(cache.mu_pred[t])). This is fine because:

• For mutable arrays (Matrix): fill_zero!! fills in-place, so the discarded return is harmless.
• For immutable arrays (SMatrix): fill_zero!! returns a new zero value (discarded), but C_t/H_t are never actually read in the immutable path of mul_aat!! (the else branch never touches them).

Since cache is a NamedTuple (immutable), cache.C_t = ... would not compile anyway, so this is the correct pattern for the mutable case.

The zeroing itself is technically unnecessary (the transpose! call overwrites the buffer before it's used as a mul! input), but it's harmless and consistent with defensive AD practice.

Correctness at the call sites

Σp = muladd_aat!!(Σp, C, cache.C_t)   # Σ̂_t = AΣA' + CC'  ✓
S  = muladd_aat!!(S,  H, cache.H_t)   # S_t = GΣ̂G' + HH'  ✓

Both call sites are correct: they replace the old muladd!!(Σp, C, transpose(C)) / muladd!!(S, H, transpose(H)) with the new workaround, computing the same result while avoiding the syrk path.

No correctness issues found. The implementation is a clean, minimal workaround for the Enzyme bug.