Fix Block matrix handling in Sum kernels and quasiseparable operations

[dfm]

Summary

This PR fixes compatibility issues with Block transition matrices in Sum kernels by adding proper handling in quasiseparable operations and introducing an optional use_block parameter to the Sum kernel class.

Key Changes

• Added use_block parameter to Sum kernel: The Sum class now accepts a use_block parameter (default True) that controls whether Block diagonal matrices or dense block_diag representations are used for transition matrices, design matrices, and stationary covariance. This provides a workaround for operations that don't yet support Block matrices.

• Fixed Block matrix handling in self_add operation: Updated StrictLowerTriQSM.self_add() to convert Block matrices to dense arrays before concatenation, preventing type errors when summing quasiseparable matrices from Sum kernels.

• Fixed Block matrix handling in self_mul operation: Updated StrictLowerTriQSM.self_mul() to handle Block matrices in the a field by converting them to dense arrays before indexing and multiplication.

• Fixed Block matrix handling in _prod_helper: Added conversion of Block matrices to dense arrays in the product kernel helper function to support products involving Sum kernels.

• Fixed Block matrix handling in qsm_mul operation: Updated the qsm_mul function in ops.py to convert Block matrices to dense arrays before concatenation operations, enabling proper conditioning with Sum kernels and banded noise.

Implementation Details

• A new _ensure_dense() helper function was added to ops.py to centralize Block-to-dense conversion logic.
• The _block_or_dense() method in the Sum class encapsulates the choice between Block and dense representations.
• All Block conversions use the .to_dense() method to maintain consistency.
• The changes are backward compatible—existing code using Sum kernels will continue to use Block matrices by default for performance, while new code can opt into dense representations if needed.

Testing

Comprehensive regression tests were added covering:

• Sum kernels with banded noise
• Sum kernels in products
• Products of Sum kernels
• Consistency between block and non-block representations
• JIT compilation and gradient computation through Sum kernels

https://claude.ai/code/session_01Y2ACGEqvh9fTrCzR5WEPuJ