Skip to content

Allow further prediction and conditioning of conditioned GPs - #287

Open
AHMETHAKANBEZIR1 wants to merge 1 commit into
dfm:mainfrom
AHMETHAKANBEZIR1:fix/static-conditioned-mean
Open

AHMETHAKANBEZIR1 wants to merge 1 commit into
dfm:mainfrom
AHMETHAKANBEZIR1:fix/static-conditioned-mean

Conversation

@AHMETHAKANBEZIR1

Copy link
Copy Markdown

A GP returned by condition cannot predict or condition on another batch at new coordinates: its Conditioned mean stores include_mean as a dynamic PyTree leaf, then evaluates that traced Boolean in a Python if. The public predict and _condition JIT paths raise TracerBoolConversionError.

Mark this configuration flag static, consistent with the public conditioning methods' treatment of include_mean. Keep all mean/covariance calculations unchanged.

Validation on Windows / Python 3.12.14 / JAX 0.12.0 / Equinox 0.13.8, CPU:

  • Four new regressions fail on unchanged main, covering both include-mean settings and both subsequent prediction and conditioning. All pass with this fix.
  • New and existing GP tests: 14 passed.
  • Full suite with default float32: 174 passed, 2 skipped; with JAX_ENABLE_X64=1: 174 passed, 2 skipped. Skips are optional george/celerite comparisons.
  • Source doctests: 3 passed. Configured pre-commit hooks and git diff --check pass; a news fragment is included.

The tests compare two-stage posterior means, predictive variances, and conditional covariances with independently assembled dense Gaussian conditioning formulas.

Prepared autonomously with OpenAI Codex, including code and tests; no independent human review is claimed.

Co-authored-by: OpenAI Codex <noreply@openai.com>

This branch has not been deployed

No deployments
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