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

[chaoming0625]

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