brainstate

Working agreement

Before writing any code, describe approach, wait for approval.
Requirements ambiguous? Ask clarifying questions before writing code.
After writing code, list edge cases + suggest test cases.
Bug? Write a test that reproduces it, then fix until the test passes.
Every correction: reflect on the mistake, plan to avoid repeating it.
All updates must be happened on the worktree branch, not main.
Use brainstate.random instead of jax.random directly for all random number generation.
Write spec and plan under doc/specs before implementation, so they're available for reference during implementation.
Tests should >90% coverage, but focus on meaningful tests that cover edge cases and critical paths, not just trivial lines.
Maintain compatibility with JAX versions >= 0.8.0. Guard any version-specific behavior in brainstate/_compatible_import.py, preferring feature/shape detection over hard version checks.

All public classes, methods, functions must use NumPy-style docstrings. Canonical section order:

Short summary – one-line imperative description (no blank line before).
Extended summary – optional, follow blank line after short summary.
Parameters – each entry: name : type on own line, description indented below.
Returns / Yields – same format as Parameters.
Raises – exception type and when raised.
See Also – related functions / classes.
Notes – implementation details, math, references.
References – numbered bibliography entries (.. [1]).
Examples – runnable, doctestable code snippets.

Wrap example code in .. code-block:: python directive so Sphinx render with syntax highlighting.
Prefix every input line with >>> (continuation lines with ...) for doctest compatibility.
Show expected output on line immediately after statement, without prompt prefix.
Separate distinct scenarios with blank >>> line.
Always include necessary imports (import brainunit as u, etc.) at top of example block so self-contained.