fix(kernels): skip non-batch buffers in Kernel.__getitem__ and expand_batch#2757
Open
umi008 wants to merge 5 commits into
Open
fix(kernels): skip non-batch buffers in Kernel.__getitem__ and expand_batch#2757umi008 wants to merge 5 commits into
umi008 wants to merge 5 commits into
Conversation
…with active_dims)
…s-gp#2591) The pre-commit require-ascii hook detected a non-ASCII em dash character (—, U+2014) in the new comment from commit 40c4a88. Pre-commit status: previously SKIPPED (not installed in env). Now passes after this fix.
umi008
force-pushed
the
fix-2591-kernel-batch-getitem
branch
from
July 10, 2026 21:02
142b06c to
b94d9d8
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #2591
Root cause
Kernel.__getitem__andKernel.expand_batchapply batch indices to all registered buffers, including non-batch metadata likeactive_dims. Whenactive_dimshas fewer elements than the batch size (e.g. shape(1,)withbatch_shape=(2,)), indexing withkernel(x)[1,:,:]raisesIndexError: index 1 is out of bounds. When the shapes happen to align, the buffer is silently corrupted to a scalar.Fix
Add a shape guard in both methods that skips buffers whose leading dimensions do not match the kernel's batch shape. Buffers that genuinely carry the batch dimensions (e.g. parameters,
randn_weights) continue to be indexed/expanded exactly as before.Changes
gpytorch/kernels/kernel.pyactive_dims/non-batch-buffer guard in__getitem__andexpand_batchtest/lazy/test_lazy_evaluated_kernel_tensor.pyScaleKernelnesting, semantic correctness,active_dims=None, direct kernel indexing, andexpand_batchscenariosTest plan
test/lazysuite (204 tests) passestest/kernelssuite (466 tests) passespre-commit run --files gpytorch/kernels/kernel.py test/lazy/test_lazy_evaluated_kernel_tensor.pypasses (all hooks, includingufmt,pyupgrade,require-ascii)