Skip to content

Support batch scales for low-rank-plus-diagonal operators - #148

Open
antonsoo wants to merge 1 commit into
cornellius-gp:mainfrom
antonsoo:fix/lowrank-batch-scaling
Open

antonsoo wants to merge 1 commit into
cornellius-gp:mainfrom
antonsoo:fix/lowrank-batch-scaling

Conversation

@antonsoo

Copy link
Copy Markdown

Multiplying a batched low-rank-plus-diagonal covariance by per-model positive scalars raises Boolean value of Tensor with more than one value is ambiguous in _mul_constant.

Use (other > 0).all(), matching the root operator's guard. All-positive batch factors keep LowRankRootAddedDiagLinearOperator and its efficient solves. Any nonpositive factor still uses the existing AddedDiagLinearOperator fallback.

Regression tests compare broadcast scaling and root/diagonal/scale gradients with dense arithmetic for one and two batch dimensions, assert the specialized type is retained, and cover a zero-factor fallback.

Validation:

  • Full suite: 5,048 tests pass on both PyTorch 2.0 CPU (Python 3.10) and PyTorch 2.14 CPU (Python 3.12).
  • Recorded Diabetes-feature covariance workloads: all 64 batch-scale cases fail before and complete afterward across CPU/CUDA, float32/64, sizes 32-128 and ranks 4/8. Values and gradients agree with dense PyTorch calculations; worst float32 relative error is 3.0e-6, float64 below 5.6e-15. Ordinary matrix/scalar controls are byte-identical across sources.
  • Eight 15-step, three-model likelihood-training runs now complete and track independent dense Adam runs (worst float32 absolute trajectory difference 1.1e-3, final-parameter relative difference 2.4e-6). This does not assert bitwise optimizer-state parity or predictive improvement.
  • Configured pre-commit hooks, full flake8 and warnings-as-errors Sphinx build pass.

This fixes the specialized low-rank-plus-diagonal implementation; the block-operator scaling in #136 does not resolve it.

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.

1 participant