dev(hansbug): fix bug of jax integration

[HansBug]

Description

Related Issue

TODO

Check List

• merge the latest version source branch/repo, and resolve all the conflicts
• pass style check
• pass all the tests

[codecov]

Codecov Report

Merging #84 (5be0032) into main (421d936) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main      #84   +/-   ##
=======================================
  Coverage   98.87%   98.87%           
=======================================
  Files          43       43           
  Lines        2762     2764    +2     
=======================================
+ Hits         2731     2733    +2     
  Misses         31       31           

Flag	Coverage Δ
unittests	98.87% <100.00%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
treevalue/tree/integration/jax.py	100.00% <100.00%> (ø)
treevalue/tree/integration/torch.py	100.00% <100.00%> (ø)

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.