Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Migrate trace_mode to linear_operators and disable conditional control flow #2163

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

feynmanliang
Copy link
Contributor

We would like to use gpytorch in our NUTS sampler. While the linked code works, it does not utilize JIT compilation (use_nnc) because of control flow.

I understand trace_mode was originally used to disable custom pytorch functions back when JIT did not support compiling through them. This PR (along with cornellius-gp/linear_operator#27) revives the setting for the AoT JIT compiler (https://pytorch.org/functorch/nightly/aot_autograd.html).

  • Moves settings.trace_mode from gpytorch to linear_operators
  • Skips conditional control flow when trace_mode is true, currently:
    • PSD check / jitter in linear_operators.utils.cholesky
    • Diagonal x1_eq_x2 check in gpytorch.kernels.kernel

@gpleiss
Copy link
Member

gpleiss commented Oct 22, 2022

We'll put this PR on hold until cornellius-gp/linear_operator#27 is merged in and we make a new LinearOperator release (this will happen soon).

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.

2 participants