Add heuristic JIT execution for JAX delta_x - #144
Open
steven-murray wants to merge 8 commits into
Open
Conversation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
Reviewer's GuideImplements a heuristic JIT execution policy for JAX PowerBox/LogNormalPowerBox delta_x, adds a cached JIT kernel with optional eager path, wires it into benchmarks to report cold vs warm timings for eager vs JIT across CPU/GPU, and expands tests to validate behavior and warnings. Sequence diagram for JAX PowerBox delta_x heuristic JIT executionsequenceDiagram
actor User
participant PowerBox
participant _delta_x_kernel
participant _delta_x_eager
participant warnings
User->>PowerBox: delta_x(key)
PowerBox->>PowerBox: _delta_x_calls += 1
PowerBox->>PowerBox: _resolve_key(key)
alt [not usejit and not _usejit_requested and _delta_x_calls > 1 and not _delta_x_warning_emitted]
PowerBox->>warnings: warn
PowerBox->>PowerBox: _delta_x_warning_emitted = True
end
alt [usejit]
PowerBox->>_delta_x_kernel: _delta_x_kernel
_delta_x_kernel-->>PowerBox: field
else [not usejit]
PowerBox->>_delta_x_eager: _delta_x_eager
_delta_x_eager-->>PowerBox: field
end
PowerBox-->>User: field
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## v1 #144 +/- ##
==========================================
+ Coverage 95.02% 95.35% +0.32%
==========================================
Files 12 12
Lines 1206 1248 +42
==========================================
+ Hits 1146 1190 +44
+ Misses 60 58 -2
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Hey - I've found 4 issues, and left some high level feedback:
- The new warning logic in
PowerBox.delta_xuseswarnings.warnbutwarningsis not imported inpowerbox.jax.powerbox, which will raise aNameErrorat runtime; add the import at the top of the module. - The
_delta_x_kernelcached JIT uses a closure overself; if any instance attributes that affect the computation (e.g.,pk,ensure_physical, geometry) are mutated after construction, the compiled kernel will not see those changes—consider documenting this expectation or guarding against post-construction mutation if that’s not intended.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The new warning logic in `PowerBox.delta_x` uses `warnings.warn` but `warnings` is not imported in `powerbox.jax.powerbox`, which will raise a `NameError` at runtime; add the import at the top of the module.
- The `_delta_x_kernel` cached JIT uses a closure over `self`; if any instance attributes that affect the computation (e.g., `pk`, `ensure_physical`, geometry) are mutated after construction, the compiled kernel will not see those changes—consider documenting this expectation or guarding against post-construction mutation if that’s not intended.
## Individual Comments
### Comment 1
<location path="src/powerbox/jax/powerbox.py" line_range="248-252" />
<code_context>
+ def delta_x(self, key: jax.Array | None = None) -> jax.Array:
+ """Return the realized real-space field using the configured execution policy."""
+ self._delta_x_calls += 1
+ if (
+ not self.usejit
+ and not self._usejit_requested
+ and self._delta_x_calls > 1
+ and not self._delta_x_warning_emitted
+ ):
+ warnings.warn(
+ "delta_x() is using eager execution by default for this box size. "
+ "Repeated calls may be much slower than usejit=True.",
+ stacklevel=2,
+ )
+ self._delta_x_warning_emitted = True
+
+ run_key = self._resolve_key(key)
</code_context>
<issue_to_address>
**suggestion (performance):** The eager-execution warning message can be misleading given the JIT heuristic.
Because the default policy prefers eager for small `Ntot` (where JIT overhead can make JIT slower), this warning can push users with small boxes and repeated calls toward `usejit=True` even when that hurts performance. Consider either:
- Emitting the warning only when `Ntot` exceeds the JIT threshold, or
- Rewording to something like "Consider usejit=True for large boxes" so it doesn’t imply JIT is generally faster in all cases.
```suggestion
warnings.warn(
"delta_x() is using eager execution by default for this box size. "
"For large boxes with repeated calls, consider usejit=True for better performance.",
stacklevel=2,
)
```
</issue_to_address>
### Comment 2
<location path="tests/test_jax.py" line_range="147-163" />
<code_context>
+ assert jit_pb.usejit is True
+
+
+def test_jax_default_eager_mode_warns_only_on_repeated_calls(monkeypatch) -> None:
+ monkeypatch.setattr(jpb_powerbox, "DEFAULT_JIT_NTOT_THRESHOLD", 10_000)
+ pb = jpb.PowerBox(
+ (8, 8),
+ dim=2,
+ pk=lambda k: (1 + k) ** -2.0,
+ boxlength=4.0,
+ key=jax.random.key(27),
+ )
+
+ with warnings.catch_warnings(record=True) as record:
+ warnings.simplefilter("always")
+ pb.delta_x()
+ pb.delta_x()
+
+ assert len(record) == 1
+ assert "Repeated calls may be much slower" in str(record[0].message)
+
+
</code_context>
<issue_to_address>
**suggestion (testing):** Add a complementary test that default JIT mode (Ntot above threshold) does not emit warnings on repeated calls.
To cover the JIT path as well, please add a counterpart test where `DEFAULT_JIT_NTOT_THRESHOLD` is set below `Ntot` so that `usejit` is `True` by default, and then assert that repeated `delta_x()` calls do not emit warnings. This will help catch regressions where warnings might appear even when JIT is enabled.
```suggestion
def test_jax_default_eager_mode_warns_only_on_repeated_calls(monkeypatch) -> None:
monkeypatch.setattr(jpb_powerbox, "DEFAULT_JIT_NTOT_THRESHOLD", 10_000)
pb = jpb.PowerBox(
(8, 8),
dim=2,
pk=lambda k: (1 + k) ** -2.0,
boxlength=4.0,
key=jax.random.key(27),
)
with warnings.catch_warnings(record=True) as record:
warnings.simplefilter("always")
pb.delta_x()
pb.delta_x()
assert len(record) == 1
assert "Repeated calls may be much slower" in str(record[0].message)
def test_jax_default_jit_mode_does_not_warn_on_repeated_calls(monkeypatch) -> None:
# Set the threshold low enough that Ntot for the box is above it,
# so that JIT is enabled by default.
monkeypatch.setattr(jpb_powerbox, "DEFAULT_JIT_NTOT_THRESHOLD", 1)
pb = jpb.PowerBox(
(8, 8),
dim=2,
pk=lambda k: (1 + k) ** -2.0,
boxlength=4.0,
key=jax.random.key(28),
)
assert pb.usejit is True
with warnings.catch_warnings(record=True) as record:
warnings.simplefilter("always")
pb.delta_x()
pb.delta_x()
# When JIT is enabled by default, repeated calls should not emit
# the "Repeated calls may be much slower" warning.
assert not any(
"Repeated calls may be much slower" in str(warning.message)
for warning in record
)
```
</issue_to_address>
### Comment 3
<location path="docs/demos/backend_benchmarks.rst" line_range="15" />
<code_context>
-shapes ``(N,)*dim`` and ``boxlength=100``. The current size ladders are
+The benchmark script is ``docs/demos/benchmark_backends.py``. It records both a
+``cold`` timing (first run, including any one-time setup such as JAX compilation) and a
+``warm`` timing (median of repeated post-warmup runs) for each configuration, using
+cubic boxes with shapes ``(N,)*dim`` and ``boxlength=100``. The current size ladders are
``[4096, 16384, 65536, 262144]`` in 1D, ``[64, 128, 256, 512]`` in 2D, and
</code_context>
<issue_to_address>
**nitpick (typo):** Consider making "post-warmup" consistent with the earlier "warm-up" spelling.
Earlier in this paragraph you use "warm-up run", but here you have "post-warmup runs". Consider changing this to "post–warm-up runs" (or "post-warm-up runs") for consistent hyphenation.
```suggestion
``warm`` timing (median of repeated post-warm-up runs) for each configuration, using
```
</issue_to_address>
### Comment 4
<location path="src/powerbox/jax/powerbox.py" line_range="221" />
<code_context>
- def delta_x(self, key: jax.Array | None = None) -> jax.Array:
- """Return the realized real-space field from the input power spectrum."""
+ def _delta_x_eager(self, key: jax.Array | None = None) -> jax.Array:
+ """Return the realized real-space field without JIT compilation."""
dk = jnp.sqrt(self._power_array_rfft()) * self._gaussian_modes_rfft(key=key)
</code_context>
<issue_to_address>
**issue (complexity):** Consider unifying eager and JIT execution behind a single `_delta_x_impl` chosen once in `__init__`, so `delta_x` becomes a simple dispatcher and subclasses only override `_delta_x_impl` rather than redefining `delta_x` or using a cached kernel.
You can keep all current behavior while flattening the execution path and clarifying the subclass contract by:
1. **Unifying eager/JIT around a single `_delta_x_impl`**
2. **Binding the chosen implementation once in `__init__`**
3. **Letting subclasses override `_delta_x_impl` only (no `delta_x` override)**
4. **Dropping the `cached_property` kernel that closes over `self`**
### 1. Base class: single impl + bound dispatcher
Replace `_delta_x_eager` + `_delta_x_kernel` with a single implementation method and a bound callable chosen in `__init__`:
```python
class PowerBox:
def __init__(..., usejit: bool | None = None):
...
self._usejit_requested = usejit is not None
self.usejit = self._resolve_usejit(usejit)
self._delta_x_calls = 0
self._delta_x_warning_emitted = False
# bind the execution function once
impl = self._delta_x_impl # instance method
if self.usejit:
# jitted wrapper, but delta_x() only sees a simple callable
self._delta_x = jax.jit(lambda run_key: impl(run_key))
else:
self._delta_x = impl
...
def _delta_x_impl(self, key: jax.Array | None = None) -> jax.Array:
"""Core delta_x implementation (eager semantics)."""
dk = jnp.sqrt(self._power_array_rfft()) * self._gaussian_modes_rfft(key=key)
field = self._irfft_to_field(dk, scale=self.V)
if self.ensure_physical:
field = jnp.clip(field, -1, jnp.inf)
return field
def delta_x(self, key: jax.Array | None = None) -> jax.Array:
"""Return the realized real-space field using the configured execution policy."""
self._delta_x_calls += 1
if (
not self.usejit
and not self._usejit_requested
and self._delta_x_calls > 1
and not self._delta_x_warning_emitted
):
warnings.warn(
"delta_x() is using eager execution by default for this box size. "
"Repeated calls may be much slower than usejit=True.",
stacklevel=2,
)
self._delta_x_warning_emitted = True
run_key = self._resolve_key(key)
return self._delta_x(run_key)
```
This keeps:
- The heuristic and explicit `usejit` override.
- The call counter and one-time warning.
- A private eager implementation (`_delta_x_impl`) still available for benchmarking.
But `delta_x` is now a simple, stateless dispatcher after the warning check, and there’s no extra indirection via a `cached_property`.
### 2. Subclass: override `_delta_x_impl` only
Have `LogNormalPowerBox` override only the implementation, not `delta_x`, and reuse the base policy logic:
```python
class LogNormalPowerBox(PowerBox):
def _delta_x_impl(self, key: jax.Array | None = None) -> jax.Array:
"""Return the realized lognormal over-density field (eager semantics)."""
dk = jnp.sqrt(self.gaussian_power_array())
dk = dk * self._gaussian_modes_rfft(key=key)
field = self._irfft_to_field(dk, scale=jnp.sqrt(self.V))
sigma_g = jnp.var(field)
return jnp.exp(field - sigma_g / 2) - 1
```
You can then drop:
- `PowerBox._delta_x_eager`
- `PowerBox._delta_x_kernel` (`@cached_property`)
- `LogNormalPowerBox._delta_x_eager`
- `LogNormalPowerBox.delta_x` override
This removes the subtle base-class contract “`delta_x` calls `_delta_x_eager`” and the closure-based cached kernel, while preserving JIT control, eager benchmarking, and the warning behavior.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
+147
to
+163
| def test_jax_default_eager_mode_warns_only_on_repeated_calls(monkeypatch) -> None: | ||
| monkeypatch.setattr(jpb_powerbox, "DEFAULT_JIT_NTOT_THRESHOLD", 10_000) | ||
| pb = jpb.PowerBox( | ||
| (8, 8), | ||
| dim=2, | ||
| pk=lambda k: (1 + k) ** -2.0, | ||
| boxlength=4.0, | ||
| key=jax.random.key(27), | ||
| ) | ||
|
|
||
| with warnings.catch_warnings(record=True) as record: | ||
| warnings.simplefilter("always") | ||
| pb.delta_x() | ||
| pb.delta_x() | ||
|
|
||
| assert len(record) == 1 | ||
| assert "Repeated calls may be much slower" in str(record[0].message) |
Contributor
There was a problem hiding this comment.
suggestion (testing): Add a complementary test that default JIT mode (Ntot above threshold) does not emit warnings on repeated calls.
To cover the JIT path as well, please add a counterpart test where DEFAULT_JIT_NTOT_THRESHOLD is set below Ntot so that usejit is True by default, and then assert that repeated delta_x() calls do not emit warnings. This will help catch regressions where warnings might appear even when JIT is enabled.
Suggested change
| def test_jax_default_eager_mode_warns_only_on_repeated_calls(monkeypatch) -> None: | |
| monkeypatch.setattr(jpb_powerbox, "DEFAULT_JIT_NTOT_THRESHOLD", 10_000) | |
| pb = jpb.PowerBox( | |
| (8, 8), | |
| dim=2, | |
| pk=lambda k: (1 + k) ** -2.0, | |
| boxlength=4.0, | |
| key=jax.random.key(27), | |
| ) | |
| with warnings.catch_warnings(record=True) as record: | |
| warnings.simplefilter("always") | |
| pb.delta_x() | |
| pb.delta_x() | |
| assert len(record) == 1 | |
| assert "Repeated calls may be much slower" in str(record[0].message) | |
| def test_jax_default_eager_mode_warns_only_on_repeated_calls(monkeypatch) -> None: | |
| monkeypatch.setattr(jpb_powerbox, "DEFAULT_JIT_NTOT_THRESHOLD", 10_000) | |
| pb = jpb.PowerBox( | |
| (8, 8), | |
| dim=2, | |
| pk=lambda k: (1 + k) ** -2.0, | |
| boxlength=4.0, | |
| key=jax.random.key(27), | |
| ) | |
| with warnings.catch_warnings(record=True) as record: | |
| warnings.simplefilter("always") | |
| pb.delta_x() | |
| pb.delta_x() | |
| assert len(record) == 1 | |
| assert "Repeated calls may be much slower" in str(record[0].message) | |
| def test_jax_default_jit_mode_does_not_warn_on_repeated_calls(monkeypatch) -> None: | |
| # Set the threshold low enough that Ntot for the box is above it, | |
| # so that JIT is enabled by default. | |
| monkeypatch.setattr(jpb_powerbox, "DEFAULT_JIT_NTOT_THRESHOLD", 1) | |
| pb = jpb.PowerBox( | |
| (8, 8), | |
| dim=2, | |
| pk=lambda k: (1 + k) ** -2.0, | |
| boxlength=4.0, | |
| key=jax.random.key(28), | |
| ) | |
| assert pb.usejit is True | |
| with warnings.catch_warnings(record=True) as record: | |
| warnings.simplefilter("always") | |
| pb.delta_x() | |
| pb.delta_x() | |
| # When JIT is enabled by default, repeated calls should not emit | |
| # the "Repeated calls may be much slower" warning. | |
| assert not any( | |
| "Repeated calls may be much slower" in str(warning.message) | |
| for warning in record | |
| ) |
|
|
||
| def delta_x(self, key: jax.Array | None = None) -> jax.Array: | ||
| """Return the realized real-space field from the input power spectrum.""" | ||
| def _delta_x_eager(self, key: jax.Array | None = None) -> jax.Array: |
Contributor
There was a problem hiding this comment.
issue (complexity): Consider unifying eager and JIT execution behind a single _delta_x_impl chosen once in __init__, so delta_x becomes a simple dispatcher and subclasses only override _delta_x_impl rather than redefining delta_x or using a cached kernel.
You can keep all current behavior while flattening the execution path and clarifying the subclass contract by:
- Unifying eager/JIT around a single
_delta_x_impl - Binding the chosen implementation once in
__init__ - Letting subclasses override
_delta_x_implonly (nodelta_xoverride) - Dropping the
cached_propertykernel that closes overself
1. Base class: single impl + bound dispatcher
Replace _delta_x_eager + _delta_x_kernel with a single implementation method and a bound callable chosen in __init__:
class PowerBox:
def __init__(..., usejit: bool | None = None):
...
self._usejit_requested = usejit is not None
self.usejit = self._resolve_usejit(usejit)
self._delta_x_calls = 0
self._delta_x_warning_emitted = False
# bind the execution function once
impl = self._delta_x_impl # instance method
if self.usejit:
# jitted wrapper, but delta_x() only sees a simple callable
self._delta_x = jax.jit(lambda run_key: impl(run_key))
else:
self._delta_x = impl
...
def _delta_x_impl(self, key: jax.Array | None = None) -> jax.Array:
"""Core delta_x implementation (eager semantics)."""
dk = jnp.sqrt(self._power_array_rfft()) * self._gaussian_modes_rfft(key=key)
field = self._irfft_to_field(dk, scale=self.V)
if self.ensure_physical:
field = jnp.clip(field, -1, jnp.inf)
return field
def delta_x(self, key: jax.Array | None = None) -> jax.Array:
"""Return the realized real-space field using the configured execution policy."""
self._delta_x_calls += 1
if (
not self.usejit
and not self._usejit_requested
and self._delta_x_calls > 1
and not self._delta_x_warning_emitted
):
warnings.warn(
"delta_x() is using eager execution by default for this box size. "
"Repeated calls may be much slower than usejit=True.",
stacklevel=2,
)
self._delta_x_warning_emitted = True
run_key = self._resolve_key(key)
return self._delta_x(run_key)This keeps:
- The heuristic and explicit
usejitoverride. - The call counter and one-time warning.
- A private eager implementation (
_delta_x_impl) still available for benchmarking.
But delta_x is now a simple, stateless dispatcher after the warning check, and there’s no extra indirection via a cached_property.
2. Subclass: override _delta_x_impl only
Have LogNormalPowerBox override only the implementation, not delta_x, and reuse the base policy logic:
class LogNormalPowerBox(PowerBox):
def _delta_x_impl(self, key: jax.Array | None = None) -> jax.Array:
"""Return the realized lognormal over-density field (eager semantics)."""
dk = jnp.sqrt(self.gaussian_power_array())
dk = dk * self._gaussian_modes_rfft(key=key)
field = self._irfft_to_field(dk, scale=jnp.sqrt(self.V))
sigma_g = jnp.var(field)
return jnp.exp(field - sigma_g / 2) - 1You can then drop:
PowerBox._delta_x_eagerPowerBox._delta_x_kernel(@cached_property)LogNormalPowerBox._delta_x_eagerLogNormalPowerBox.delta_xoverride
This removes the subtle base-class contract “delta_x calls _delta_x_eager” and the closure-based cached kernel, while preserving JIT control, eager benchmarking, and the warning behavior.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
usejitpolicy forpowerbox.jax.PowerBox/LogNormalPowerBox.delta_x()use a cached JIT kernel by default, with a private eager path for comparison and benchmarks.get_powerusage andusejitguidance.Target
v1Summary by Sourcery
Introduce a heuristic-controlled JIT execution path for JAX PowerBox real-space field generation and extend benchmarks to compare eager vs JIT performance with cold and warm timings.
New Features:
Enhancements:
Tests: