Repository navigation
fix: High/Medium audit bugs across state, environ, random, nn, transform - #244
Merged
Merged
Conversation
…mpare Mode parameters - Mode.__eq__/__hash__ include Batching(batch_size, batch_axis) and JointMode components, and return NotImplemented for non-Mode operands, so jit static Mode arguments retrace when parameters change. - update_from_ref keeps the state's own HookManager and source info; TreefyState.to_state and State.replace get a fresh HookManager; copy, replace, copy_from and to_state copy the tag set instead of aliasing it. - Hook.execute re-raises HookCancellationError so raising it cancels a write; logged hook errors report the original exception. - HookManager rebuilds its cache under the lock and runs callbacks outside it. - StateTraceStack.__exit__ no longer masks an in-flight exception and removes itself (and leaked entries above it) from the stack.
…s threads - EnvironmentState settings and callbacks are shared by all threads; only the context() override stack stays per thread, so worker threads see dt/precision. - Overlapping context(precision=...) blocks across threads restore the most recent still-active override instead of clobbering the global JAX flag; with none left, the global precision setting (or the import-time precision) wins.
…s, zipf size - pareto subtracts 1 from jax's Pareto-I samples so support starts at 0, as in numpy and the docstring. - multinomial treats the last pval as 1 - sum(rest) like numpy and clips the sampled index, so counts always sum to n. - zipf normalizes size via _size2shape (list sizes no longer fail as unhashable static args).
…ing, PrefetchDelay - BatchNorm with track_running_stats=False normalizes with batch statistics instead of returning the input unchanged. - BatchNorm reshapes batch mean/var to the running-stat shape before the EMA update, so feature_axis != -1 no longer corrupts running stats. - MaxUnpool scatters indices against the original input layout when output_size matches it (fixes odd-sized inputs); MaxPool(return_indices=True, in_size=...) no longer crashes in the constructor. - Delay rounds float delay/dt ratios that are within 1e-6 of an integer instead of ceil-ing them, and update_every_step uses round() (0.3/0.1 -> 3). - init_maybe_prefetch(PrefetchDelay) resolves the registered delay instead of raising TypeError.
…onate argnames - DCE keeps equations with effects (debug_print, callbacks, jit_error_if) and their inputs; constant folding skips effectful equations, including nested jit. - x - x and 0 * x simplifications only apply to 0-d integer outputs (no wrong shape, NaN preserved). - vmap_new_states / vmap2_new_states roll back pre-existing states written with mapped tracers and raise a clear error; states excluded via state_to_exclude are poisoned and unwound instead of keeping dead tracers. - StatefulFunction / make_jaxpr normalize negative static_argnums per call. - jit infers static/donate argnums from argnames (and vice versa) via the function signature, like jax.jit. - Mixed batched/unbatched leaves in one state are accepted (unbatched leaves broadcast). - Python code generation for scan with consts, const-only body ops and unused nested-jaxpr invars produces runnable code.
…add multi-device tests jax.pmap does not invalidate its inner DynamicJaxprTrace, so the leaked-state check of pmap2_new_states missed it. A tracer now also counts as leaked when its trace is not on the current trace's parent_trace chain. Tests: pmap2 with mixed batched/unbatched state leaves, jit static_argnames passed positionally under pmap2, pmap2_new_states rollback, per-device callbacks, and outer-jit tracer writes that must not be flagged.
… against callback deadlocks
…size cases; document output_size
…t_size MaxPool(return_indices=True) returned flat positions in the original input, whose size MaxUnpool cannot recover (4 and 5 both pool to 2), so odd-sized inputs were silently misplaced without output_size, even across batch elements. Indices are now flat positions in a layout of (out - 1) * stride + kernel - pad_lo per pooled axis, which depends only on the pooled shape; MaxUnpool unravels against it and re-ravels into the natural size or the requested output_size. Values are unchanged for evenly divisible, unpadded inputs.
Contributor
|
Sorry @chaoming0625, you've used your own review budget of 250,000 diff characters for the last 7 days. You can request another review in 6 days and 1 hour by commenting |
Contributor
Reviewer's GuideThis PR implements a package-wide High/Medium audit remediation: it isolates state and hook metadata, removes hook and environment concurrency hazards, corrects random and neural-network edge-case semantics, and makes transform optimizations and stateful mappings preserve effects and tracer safety, with broad regression coverage and simulated multi-device testing. Sequence diagram for deadlock-safe state hook executionsequenceDiagram
participant State
participant HookManager
participant Lock
participant Hook
State->>HookManager: execute_write_before_hooks()
HookManager->>Lock: _get_cached_hooks()
Lock-->>HookManager: fresh hook snapshot
HookManager->>Lock: release
HookManager->>Hook: execute(context)
Hook-->>HookManager: transformed value or HookCancellationError
alt hook cancels write
HookManager-->>State: raise HookCancellationError
State-->>State: cancel write
else hook succeeds
HookManager-->>State: transformed value
end
Sequence diagram for tracer-safe mapped state updatessequenceDiagram
participant Transform
participant MappedFunction
participant TraceStack
participant State
participant Cleanup
Transform->>TraceStack: begin mapped pass
Transform->>MappedFunction: vmap_new_states() or pmap2_new_states()
MappedFunction->>State: write mapped tracer
MappedFunction-->>Transform: mapped result
Transform->>Cleanup: rollback_leaked_preexisting_states()
Cleanup->>State: restore original value
alt pre-existing state received mapped tracer
Cleanup-->>Transform: leaked states
Transform->>Transform: raise_leaked_preexisting_states()
else newly created state
Cleanup->>TraceStack: unwind_new_state_levels()
Transform-->>Transform: continue with mapped states
end
Flow diagram for size-independent MaxPool and MaxUnpool indicesflowchart LR
Input[Odd or padded input] --> Pool[MaxPool return_indices]
Pool --> Layout[Size-independent index layout]
Pool --> Values[Pooled values]
Layout --> Unpool[MaxUnpool]
Values --> Unpool
OutputSize[Optional output_size] --> Unpool
Unpool --> Output[Correctly scattered output]
Flow diagram for side-effect-preserving transform optimizationflowchart LR
Jaxpr[Jaxpr equations] --> Inspect[Inspect outputs and effects]
Inspect -->|side effects present| Keep[Keep equation at runtime]
Inspect -->|no effects and constant inputs| Fold[Constant fold]
Inspect -->|no effects and unused outputs| Remove[Remove equation]
Keep --> Optimized[Optimized Jaxpr]
Fold --> Optimized
Remove --> Optimized
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Merged
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
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
Fixes the High/Medium bugs found in a package-wide audit. Every fix has a regression test that fails before the change. The spec is in
doc/specs/2026-10-audit-fixes.md.Core state & hooks
Modeequality and hashing now include parameters (Batching(batch_size, batch_axis)and the components ofJointMode), sojitretraces when a static mode changes. Comparing with a non-Mode returnsNotImplemented.update_from_ref,replace,copy,copy_fromandTreefyState.to_stateno longer share or lose hook managers or tag sets.HookCancellationErrorinside a hook now cancels the write. Logged hook errors report the original exception.HookManagerrebuilds its cache under the lock and runs callbacks outside it, so they can no longer deadlock.StateTraceStack.__exit__no longer hides an in-flight exception, and no longer leaves stale entries on the stack.environ
context()stack stays per thread.context(precision=...)blocks that overlap across threads no longer turn x64 off for each other.random
paretonow samples Lomax, matching numpy and its docs. Behavior change.multinomialcounts always sum ton.zipfaccepts a listsize.nn
track_running_stats=Falsenow normalizes. Running stats keep their shape whenfeature_axisis not last.MaxUnpoolhandles odd sizes withoutoutput_size.MaxPool(return_indices=True, in_size=...)no longer crashes in the constructor.update_every_stepare rounded correctly (e.g.0.3/0.1 → 3).init_maybe_prefetch(PrefetchDelay)works.transform
debug_print, callbacks,jit_error_if).x - x/0 * xrewrites now apply only to integer scalars.vmap_new_states,vmap2_new_statesandpmap2_new_statesroll back pre-existing states that were written with mapped tracers, then raise a clear error. pmap traces are detected through theparent_tracechain. States excluded viastate_to_excludeare poisoned and unwound.static_argnumswork.jitinfers static/donate argnums from argnames, and argnames from argnums, likejax.jit.scannow produces runnable code.Tests
conftest.pysimulates 8 CPU devices viajax.config.update('jax_num_cpu_devices', 8), so multi-device tests run without accelerators.Behavior changes
random.paretosamples are 1 lower than before.vmap*_new_statesraises when a batched value is written into a pre-existing state. Previously the tracer leaked silently.x.ravel()[indices]no longer gathers the maxima.Test plan
pytest brainstate -n 8: 5438 passed, 53 skipped (slow tests and missing optional frameworks), 0 failed.Summary by Sourcery
Resolve high- and medium-severity correctness, concurrency, numerical, and transformation bugs across the package while adding comprehensive regression coverage.
Bug Fixes:
Enhancements:
CI:
Documentation:
Tests: