From 399fa14407b49f4f53673747a99207cab3d2c6d4 Mon Sep 17 00:00:00 2001 From: Mark Fortune Date: Fri, 10 Apr 2026 14:12:32 +1000 Subject: [PATCH 1/4] Avoid nesting Block objects for quasisep kernels Fixing a bug which caused the addition of three or more quasiseparable kernels to fail in certain cases due to nested Block objects --- src/tinygp/kernels/quasisep.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/tinygp/kernels/quasisep.py b/src/tinygp/kernels/quasisep.py index b380f08..f2166a3 100644 --- a/src/tinygp/kernels/quasisep.py +++ b/src/tinygp/kernels/quasisep.py @@ -243,9 +243,14 @@ def coord_to_sortable(self, X: JAXArray) -> JAXArray: return self.kernel1.coord_to_sortable(X) def _block_or_dense(self, m1: JAXArray, m2: JAXArray) -> JAXArray: - if self.use_block: - return Block(m1, m2) - return jsp_block_diag(m1, m2) + if not self.use_block: + return jsp_block_diag(m1, m2) + + # Ensure we don't nest Block objects to fix Issue #265 + blocks1 = m1.blocks if isinstance(m1, Block) else (m1,) + blocks2 = m2.blocks if isinstance(m2, Block) else (m2,) + return Block(*blocks1, *blocks2) + def design_matrix(self) -> JAXArray: return self._block_or_dense( From 59e6bc6836f529accd0e9be7ee486599da8a048f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 10 Apr 2026 04:26:15 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/tinygp/kernels/quasisep.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/tinygp/kernels/quasisep.py b/src/tinygp/kernels/quasisep.py index f2166a3..2ed03f6 100644 --- a/src/tinygp/kernels/quasisep.py +++ b/src/tinygp/kernels/quasisep.py @@ -245,12 +245,11 @@ def coord_to_sortable(self, X: JAXArray) -> JAXArray: def _block_or_dense(self, m1: JAXArray, m2: JAXArray) -> JAXArray: if not self.use_block: return jsp_block_diag(m1, m2) - + # Ensure we don't nest Block objects to fix Issue #265 blocks1 = m1.blocks if isinstance(m1, Block) else (m1,) blocks2 = m2.blocks if isinstance(m2, Block) else (m2,) return Block(*blocks1, *blocks2) - def design_matrix(self) -> JAXArray: return self._block_or_dense( From 37e2fcd9193af34e478821e65cbea5d9c364c0a7 Mon Sep 17 00:00:00 2001 From: Mark Fortune Date: Tue, 14 Apr 2026 11:38:06 +1000 Subject: [PATCH 3/4] Adding block multiplication on the right Avoid crashing from scalar multiplication of kernels with block stationary covariance matrices --- src/tinygp/solvers/quasisep/block.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/tinygp/solvers/quasisep/block.py b/src/tinygp/solvers/quasisep/block.py index 6e102cd..7da2389 100644 --- a/src/tinygp/solvers/quasisep/block.py +++ b/src/tinygp/solvers/quasisep/block.py @@ -55,6 +55,9 @@ def to_dense(self) -> JAXArray: def __mul__(self, other: Any) -> "Block": return Block(*(b * other for b in self.blocks)) + def __rmul__(self, other: Any) -> "Block": + return self.__mul__(other) + @jax.jit def __add__(self, other: Any) -> Any: if isinstance(other, Block): From 9e049d37586a31431ad981486543e63a16874bc8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 14 Apr 2026 01:40:49 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/tinygp/solvers/quasisep/block.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tinygp/solvers/quasisep/block.py b/src/tinygp/solvers/quasisep/block.py index 7da2389..3bcb0bc 100644 --- a/src/tinygp/solvers/quasisep/block.py +++ b/src/tinygp/solvers/quasisep/block.py @@ -57,7 +57,7 @@ def __mul__(self, other: Any) -> "Block": def __rmul__(self, other: Any) -> "Block": return self.__mul__(other) - + @jax.jit def __add__(self, other: Any) -> Any: if isinstance(other, Block):