Jax windows restrictions

[kratman]

Description

Restrictions on testing and installing Jax can be removed after #3550

Related: #3941

Key checklist:

• No style issues: $ pre-commit run (or $ nox -s pre-commit) (see CONTRIBUTING.md for how to set this up to run automatically when committing locally, in just two lines of code)
• All tests pass: $ python run-tests.py --all (or $ nox -s tests)
• The documentation builds: $ python run-tests.py --doctest (or $ nox -s doctests)

You can run integration tests, unit tests, and doctests together at once, using $ python run-tests.py --quick (or $ nox -s quick).

Further checks:

• Code is commented, particularly in hard-to-understand areas
• Tests added that prove fix is effective or that feature works

[kratman]

@agriyakhetarpal This should be all of the jax stuff. Since I am waiting on #3954 I will do the coverage fixes in a third PR

[agriyakhetarpal]

Thanks!

[kratman]

Sorry pushed another change after your review

dismissing my earlier review

[kratman]

Coverage is still failing in #3954, so it will fail here too

[agriyakhetarpal]

I updated the branch, it should be good to go now. Please feel free to merge

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.58%. Comparing base (78b7d25) to head (685ea0f).

Additional details and impacted files

@@           Coverage Diff            @@
##           develop    #3955   +/-   ##
========================================
  Coverage    99.58%   99.58%           
========================================
  Files          257      257           
  Lines        21196    21196           
========================================
  Hits         21108    21108           
  Misses          88       88           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

There are some failing checks

[agriyakhetarpal]

Based on the failing tests, I think the long-standing issue of Homebrew's gcc not coming with gfortran is still not fixed. We should add them back again because we wouldn't want workflows like publish_pypi.yml to fail for unrelated reasons on releases.

We can remove OpenBLAS, though. It's not being used currently because we use Apple's Accelerate framework.

[agriyakhetarpal]

Had to trigger a re-run for some, but all tests pass now, we should be good to go. Thanks!