build(sdk): temporarily pin jax in tests because of breaking changes

[raubitsj]

Description

jax 0.4.1 changed something which is breaking our tests... temporarily pin so we can keep tests passing while we figure it out

Testing

circleci

[codecov]

Codecov Report

Merging #4632 (a2fe3cf) into main (39d50c3) will decrease coverage by 0.01%.
The diff coverage is n/a.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4632      +/-   ##
==========================================
- Coverage   83.22%   83.20%   -0.02%     
==========================================
  Files         279      279              
  Lines       34903    34903              
==========================================
- Hits        29047    29041       -6     
- Misses       5856     5862       +6     

Flag	Coverage Δ
functest	56.10% <ø> (+<0.01%)	⬆️
unittest	73.75% <ø> (+<0.01%)	⬆️

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

Impacted Files	Coverage Δ
wandb/integration/tensorboard/monkeypatch.py	91.35% <0.00%> (-2.47%)	⬇️
wandb/sdk/lib/sock_client.py	93.12% <0.00%> (-1.06%)	⬇️
wandb/sdk/lib/printer.py	88.60% <0.00%> (-0.64%)	⬇️
wandb/sdk/lib/git.py	77.71% <0.00%> (-0.61%)	⬇️
wandb/util.py	87.29% <0.00%> (-0.41%)	⬇️
wandb/sdk/wandb_run.py	90.60% <0.00%> (-0.06%)	⬇️
wandb/sdk/internal/system/system_info.py	91.55% <0.00%> (+3.24%)	⬆️

[kptkin]

approve

[philipjball]

I believe this might be related to the recent API changes here, where everything is now a jax.Array, subsuming jax.DeviceArray.

In any case, Jax 0.4.1 breaks logging because Jax arrays no longer pass the typename check in util.py.

[dmitryduev]

thanks @philipjball, will fix!

[dmitryduev]

hey @philipjball, the fix should be in main -- would you mind checking that it works for you with jax==0.4.1?

[philipjball]

Yep works perfectly for me! Thanks for fixing that so quickly :) Below is just a mockup of two jax Arrays I've logged successfully using 0.4.1: