Use jax key registry

[ASEM000]

• bump min jax version to 0.4.7 to use the lastest jax key registry to replace the internal registry.
• tree_map_with_trace accept trace of a tuple of jax key entry, type
• use jax.tree_util.keystr.

[review-notebook-app[bot]]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov[bot]]

Codecov Report

Patch coverage: 97.84% and project coverage change: -0.11 :warning:

Comparison is base (9814ca1) 99.63% compared to head (4ee8e03) 99.53%.

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #58 +/- ## ========================================== - Coverage 99.63% 99.53% -0.11% ========================================== Files 14 14 Lines 2493 2356 -137 ========================================== - Hits 2484 2345 -139 - Misses 9 11 +2 ``` | [Impacted Files](https://codecov.io/gh/ASEM000/PyTreeClass/pull/58?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Mahmoud+Asem) | Coverage Δ | | |---|---|---| | [pytreeclass/\_\_init\_\_.py](https://codecov.io/gh/ASEM000/PyTreeClass/pull/58?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Mahmoud+Asem#diff-cHl0cmVlY2xhc3MvX19pbml0X18ucHk=) | `100.00% <ø> (ø)` | | | [tests/test\_treeclass.py](https://codecov.io/gh/ASEM000/PyTreeClass/pull/58?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Mahmoud+Asem#diff-dGVzdHMvdGVzdF90cmVlY2xhc3MucHk=) | `99.10% <ø> (-0.18%)` | :arrow_down: | | [pytreeclass/\_src/tree\_indexer.py](https://codecov.io/gh/ASEM000/PyTreeClass/pull/58?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Mahmoud+Asem#diff-cHl0cmVlY2xhc3MvX3NyYy90cmVlX2luZGV4ZXIucHk=) | `99.09% <88.23%> (-0.91%)` | :arrow_down: | | [pytreeclass/\_src/tree\_decorator.py](https://codecov.io/gh/ASEM000/PyTreeClass/pull/58?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Mahmoud+Asem#diff-cHl0cmVlY2xhc3MvX3NyYy90cmVlX2RlY29yYXRvci5weQ==) | `100.00% <100.00%> (ø)` | | | [pytreeclass/\_src/tree\_pprint.py](https://codecov.io/gh/ASEM000/PyTreeClass/pull/58?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Mahmoud+Asem#diff-cHl0cmVlY2xhc3MvX3NyYy90cmVlX3BwcmludC5weQ==) | `100.00% <100.00%> (ø)` | | | [pytreeclass/\_src/tree\_trace.py](https://codecov.io/gh/ASEM000/PyTreeClass/pull/58?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Mahmoud+Asem#diff-cHl0cmVlY2xhc3MvX3NyYy90cmVlX3RyYWNlLnB5) | `100.00% <100.00%> (ø)` | | | [tests/test\_indexing.py](https://codecov.io/gh/ASEM000/PyTreeClass/pull/58?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Mahmoud+Asem#diff-dGVzdHMvdGVzdF9pbmRleGluZy5weQ==) | `99.76% <100.00%> (+<0.01%)` | :arrow_up: | | [tests/test\_tree\_pprint.py](https://codecov.io/gh/ASEM000/PyTreeClass/pull/58?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Mahmoud+Asem#diff-dGVzdHMvdGVzdF90cmVlX3BwcmludC5weQ==) | `98.33% <100.00%> (ø)` | | Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Mahmoud+Asem). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Mahmoud+Asem)

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Do you have feedback about the report comment? Let us know in this issue.