Refresh docs notebooks; fix clear_buffer/pmap runtime bugs; drive package-wide mypy to zero#863
Merged
Merged
Conversation
Re-executed every notebook under docs/ end-to-end against BrainPy 2.8.0 and refreshed their stored outputs in place (embedded figures included). Repaired notebook-level API drift where cells had gone stale: - interoperation: bp.interop.FromFlax -> bp.dnn.FromFlax, bp.layers.GRUCell -> bp.dyn.GRUCell, generator-based BPTT.fit(...), and bm.set(mode=bm.training_mode) before model construction. - parallel_for_parameter_exploration: set the XLA host-device flag before importing brainpy (import locks the backend), and pass progress_bar=False to the multi-device pmap DSRunner (ordered io_callback effects cannot cross a multi-device pmap). - gotchas_of_brainpy_transforms: fold the now-rejected Variable-as-arg call into the notebook's own try/except demonstration. - Dedicated_Operators: materialize sparse nonzero indices to numpy. All code cells run cleanly with sensible outputs.
…to zero Runtime fixes (surfaced while re-executing the docs notebooks): - math/environment.py: clear_buffer_memory() now resets JAX's runtime effect-token registry after deleting live device buffers, so a later ordered io_callback no longer crashes with "deleted or donated buffer". - running/jax_multiprocessing.py: jax_parallelize_map now routes through brainstate.transform.pmap2 (state-aware) instead of raw jax.pmap, and pulls each chunk's outputs to host before concatenation. This fixes a TraceContextError when the mapped function builds a stateful model (random init writes to a State) and a device-placement failure on the trailing partial chunk sharded on a subset of devices. Added co-located regression tests for both modules. Typing: - Drove `mypy brainpy/` from 1568 errors to 0 across 123 files. These are annotation-only changes with no runtime behavior change: implicit-Optional defaults promoted to Optional[T], missing var-annotations added, Optional values narrowed with asserts/guards, TypeVar-as-attribute annotations widened, and a small number of targeted `# type: ignore[code]` where the offending annotation lives in a third-party package. The curated CI mypy subset remains green.
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
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
Three related pieces of housekeeping, done together on this branch:
1. Docs notebook refresh
Re-executed every notebook under
docs/end-to-end against BrainPy 2.8.0 and refreshed their stored outputs in place (figures embedded). Repaired notebook-level API drift where cells had gone stale:bp.interop.FromFlax→bp.dnn.FromFlax,bp.layers.GRUCell→bp.dyn.GRUCell, generator-basedBPTT.fit(...), andbm.set(mode=bm.training_mode)before model construction.progress_bar=Falseto the multi-devicepmapDSRunner.2. Runtime bugfixes (surfaced while re-executing the notebooks)
math/environment.py:clear_buffer_memory()now resets JAX's runtime effect-token registry after deleting live device buffers, so a subsequent orderedio_callbackno longer crashes with "deleted or donated buffer".running/jax_multiprocessing.py:jax_parallelize_mapnow routes throughbrainstate.transform.pmap2(state-aware) instead of rawjax.pmap, and pulls each chunk's outputs to host before concatenation. Fixes aTraceContextErrorwhen the mapped function builds a stateful model (random init writes to aState) and a device-placement failure on a trailing partial chunk sharded on a subset of devices.3. Package-wide mypy typing
Drove
mypy brainpy/from 1568 errors → 0 across 123 files. Annotation-only, no runtime behavior change: implicit-Optional defaults promoted toOptional[T], missing var-annotations added, Optional values narrowed with asserts/guards,TypeVar-as-attribute annotations widened, and a small number of targeted# type: ignore[code]where the offending annotation lives in a third-party package. The curated CI mypy subset remains green.Validation
python -m mypy brainpy/→Success: no issues found in 432 source files.Latent bugs found but NOT fixed (out of scope — flagging for follow-up)
brainpy/mixin.py:SupportSTDP.stdp_updatehas a keyword typoonn_post(should beon_post) — this is the root of the STDP[override]signature mismatches.brainpy/runners.py:455:raise warnings.warn(...)raises aTypeErrorinstead of emitting a warning.