Skip to content

fix(kernels): skip non-batch buffers in Kernel.__getitem__ and expand_batch#2757

Open
umi008 wants to merge 5 commits into
cornellius-gp:mainfrom
umi008:fix-2591-kernel-batch-getitem
Open

fix(kernels): skip non-batch buffers in Kernel.__getitem__ and expand_batch#2757
umi008 wants to merge 5 commits into
cornellius-gp:mainfrom
umi008:fix-2591-kernel-batch-getitem

Conversation

@umi008

@umi008 umi008 commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

Closes #2591

Root cause

Kernel.__getitem__ and Kernel.expand_batch apply batch indices to all registered buffers, including non-batch metadata like active_dims. When active_dims has fewer elements than the batch size (e.g. shape (1,) with batch_shape=(2,)), indexing with kernel(x)[1,:,:] raises IndexError: 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.

for buffr_name, buffr in self.named_buffers(recurse=False):
    batch_ndim = len(self.batch_shape)
    if batch_ndim > 0 and buffr.shape[:batch_ndim] != self.batch_shape:
        continue
    ...

Changes

File Change
gpytorch/kernels/kernel.py Add active_dims/non-batch-buffer guard in __getitem__ and expand_batch
test/lazy/test_lazy_evaluated_kernel_tensor.py 10 regression tests covering scalar index, slice, ScaleKernel nesting, semantic correctness, active_dims=None, direct kernel indexing, and expand_batch scenarios

Test plan

  • 10 new regression tests pass (TDD: each was failing before the fix)
  • Full test/lazy suite (204 tests) passes
  • Full test/kernels suite (466 tests) passes
  • pre-commit run --files gpytorch/kernels/kernel.py test/lazy/test_lazy_evaluated_kernel_tensor.py passes (all hooks, including ufmt, pyupgrade, require-ascii)

@umi008
umi008 force-pushed the fix-2591-kernel-batch-getitem branch from 142b06c to b94d9d8 Compare July 10, 2026 21:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] Indexing LazyEvaluatedKernelTensor on batch dimension with index greater than 0 gives IndexError

1 participant