Fix for new jax version.

[janosg]

Problem

The type of a jax array has changed in jax version 0.4.1 and thus pybaum would not work with this jax version.

Solution

Handle jax arrays explicitly in get_type, as we are currently already doing for namedtuples.

[codecov]

Codecov Report

Merging #22 (d148969) into main (aa0e218) will decrease coverage by 0.28%.
The diff coverage is 91.66%.

@@            Coverage Diff             @@
##             main      #22      +/-   ##
==========================================
- Coverage   94.90%   94.61%   -0.29%     
==========================================
  Files          10       10              
  Lines         451      464      +13     
==========================================
+ Hits          428      439      +11     
- Misses         23       25       +2     

Impacted Files	Coverage Δ
src/pybaum/registry_entries.py	88.15% <ø> (-0.31%)	⬇️
src/pybaum/typecheck.py	91.30% <90.47%> (-8.70%)	⬇️
src/pybaum/equality.py	100.00% <100.00%> (ø)
tests/test_typecheck.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.