Enable brainpy object as pytree so that it can be applied with ``jax.jit`` etc. directly - Githubissues

brainpy / BrainPy

Brain Dynamics Programming in Python

https://brainpy.readthedocs.io/

GNU General Public License v3.0

515 stars 93 forks source link

Enable brainpy object as pytree so that it can be applied with ``jax.jit`` etc. directly #625

Closed chaoming0625 closed 7 months ago

chaoming0625 commented 7 months ago

Description

How Has This Been Tested

Types of changes

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Documentation (non-breaking change which updates documentation)
Breaking change (fix or feature that would cause existing functionality to change)
Code style (formatting, renaming)
Refactoring (no functional changes, no api changes)
Other (please describe here):

Checklist

[ ] Code follows the code style of this project.
[ ] Changes follow the CONTRIBUTING guidelines.
[ ] Update necessary documentation accordingly.
[ ] Lint and tests pass locally with the changes.
[ ] Check issues and pull requests first. You don't want to duplicate effort.

Other information