[Continued] Refactor botorch/sampling/pathwise and add support for product kernels

[seashoo]

Motivation

Hello! I'm Sahran Ashoor, an undergraduate research assistant working for the Uncertainty Quantification Lab at the University of Houston. I work under Dr. Ruda Zhang and Taiwo Adebiyi, both of whom having already spoken with Max Balandat regarding incorporating a rebase of botorch/sampling/pathwise (Largely written by James T. Wilson). The changes included in this pull request are my best attempt at faithfully completing the change logs provided.

Note: This PR is a continuation of PR #2838, which became inaccessible due to a branch reset issue. The original discussion (107 comments) can be found there for additional context.

Project Overview

The primary goal was to make the original codebase by Wilson compatible with the latest BoTorch version. To achieve this, we used the original source codes and test suites, which initially revealed several incompatibility issues. Our main contribution involved carefully rebasing Wilson's code while preserving the logic for pathwise sampling. Importantly, all core changes were confined to the botorch/sampling/pathwise directory to ensure a seamless integration, passing both local pathwise test suites and BoTorch's global test suites via GitHub workflows.

In terms of code logic, we relied on Wilson's unit tests for prior, updates, and posterior sampling, which we believe are sufficient to validate the correctness of the implementation. However, we welcome your feedback on this approach, and would appreciate any suggestions for additional tests or example scripts to further confirm the robustness of the changes. We are open to collaborating further on this effort.

Summary of Changes

• Refactoring botorch/sampling/pathwise for product kernel support
• Comprehensive test coverage improvements for pathwise sampling
• ProductKernel MTGP adjustments
• Compatibility updates for latest BoTorch version

Related PRs

• Original PR: #2838 (closed due to branch reset issue - contains full review discussion)

[seashoo]

@Balandat @hvarfner Happy New Year! I hope the holidays have been well for the team at Meta.

I've opened a ticket with GitHub regarding the state of the previous PR, of which I've isolated the issue to be server-side. I'll provide updates if reverting the PR to its prior state is found to be possible. In the meantime, this new PR should be almost merge-ready and includes upstream-dependent changes + lengthy testing coverage. Please feel free to provide feedback!

[hvarfner]

@seashoo Thanks! So IIUC, this is the same PR re-opened, right? Are there any changes at all to the relevant code? We would need the PR to contain only the changes in the relevant files, and then we'll take a final look.

Once again, fantastic work!

[seashoo]

Hello @hvarfner! Thank you so much for your patience on this.

With GitHub support confirming that head from the old PR is currently inaccessible, and that a continuation PR such as this is the best route to continue this project, we can continue review from this point. The most recent commit that was accessible was d85a6cc, which contained a largely intact seashoo:main containing the majority of implemented functionality that we've been seeing together throughout the year. The main divergence in the later commits were changes made in respect to the upstream of the current meta-pytorch:main, along with slightly differing approaches to test coverage. I recommend giving it another review, but I would say that all the important rebase functionality still remains.

If I'm not mistaken, models/utils/helpers.py was created as we wanted to consolidate general model utilities like get_train_inputs and get_train_targets into it's own directory. In that case, all changes seen in this PR should be directly relevant to the rebase. Please don't hesitate if you have any more questions!