Skip to content

Support traced constants on the left of kernel sums - #285

Open
AHMETHAKANBEZIR1 wants to merge 1 commit into
dfm:mainfrom
AHMETHAKANBEZIR1:fix/traced-kernel-radd
Open

AHMETHAKANBEZIR1 wants to merge 1 commit into
dfm:mainfrom
AHMETHAKANBEZIR1:fix/traced-kernel-radd

Conversation

@AHMETHAKANBEZIR1

Copy link
Copy Markdown

A traced scalar on the left of a kernel sum, such as offset + kernels.ExpSquared(...), raises TracerBoolConversionError under JIT or differentiation. The reflected-addition method tests offset == 0 in Python to support sum, before it builds the constant kernel.

Apply that shortcut only to builtin numeric zeros. Traced scalars follow the existing Constant/Sum path, while 0 + kernel, 0.0 + kernel, and sum([kernel]) keep their identity behavior.

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

  • On unchanged main: 3 failing regressions / 1 passing zero-sum control.
  • New and existing base-kernel tests: 22 passed.
  • Full test suite with default float32: 174 passed, 2 skipped; with JAX_ENABLE_X64=1: 174 passed, 2 skipped. Skips are the optional george/celerite comparisons.
  • Source doctests, with the declared doctest extra: 3 passed.
  • Configured pre-commit hooks and git diff --check pass; a news fragment is included.

Regressions check JIT, the analytic all-ones kernel derivative, vmap, and GP marginal likelihood/gradient against an independently formed dense covariance reference. No change to quasiseparable kernels' constant-addition restrictions is proposed.

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