Bump jax to 0.7.1

[rniczh]

Context:

Description of the Change:

1. API change of lower_jaxpr_to_fun, update it in catalyst accordingly
2. API change of jaxpr_subcomp, update it in catalyst accordingly
3. _jaxpr_replicas has been moved to jax._src.interpreters.pxla in jax
4. After jax 0.7.1+ mlir.multi_broadcast_in_dim has changed, and break the dynamic shape, patch it to fix the dynamic shape issue

Benefits:

Possible Drawbacks:

Related GitHub Issues:

[github-actions[bot]]

Hello. You may have forgotten to update the changelog! Please edit doc/releases/changelog-dev.md on your branch with:

• A one-to-two sentence description of the change. You may include a small working example for new features.
• A link back to this PR.
• Your name (or GitHub username) in the contributors section.

[JerryChen97]

Mind push an empty commit to trigger CI? @rniczh

[codecov[bot]]

Codecov Report

:white_check_mark: All modified and coverable lines are covered by tests. :white_check_mark: Project coverage is 97.52%. Comparing base (e0e7659) to head (c4fd765).

Additional details and impacted files

@@                     Coverage Diff                      @@
##           rniczh/bump-jax-to-0.7.0    #2134      +/-   ##
============================================================
- Coverage                     97.52%   97.52%   -0.01%     
============================================================
  Files                            93       93              
  Lines                         11080    11106      +26     
  Branches                       1058     1061       +3     
============================================================
+ Hits                          10806    10831      +25     
  Misses                          211      211              
- Partials                         63       64       +1     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

:rocket: New features to boost your workflow:

• :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.