Bump to jax 0.9.1

[dylanhmorris]

• Increments version
• Addresses warnings and failing tests seen in Bump jax from 0.9.0.1 to 0.9.1 #735

[dylanhmorris]

Todo: MRE and issue in numpyro for dist.Beta producing the following deprecation warning:

DeprecationWarning: Implicit conversion of an array to a dtype is deprecated; rather than dtype=arr use dtype=arr.dtype. In the future this will result in an error.
    key = np.dtype(dtype)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 98.31%. Comparing base (ef83fbb) to head (ad32148).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #749   +/-   ##
=======================================
  Coverage   98.31%   98.31%           
=======================================
  Files          48       48           
  Lines        1485     1486    +1     
=======================================
+ Hits         1460     1461    +1     
  Misses         25       25           

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

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

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

🚀 New features to boost your workflow:

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

[dylanhmorris]

Remaining errors will be addressed as part of #746 so this PR will wait for review until the PR closing #746 lands.

[dylanhmorris]

Todo: MRE and issue in numpyro for dist.Beta producing the following deprecation warning:

DeprecationWarning: Implicit conversion of an array to a dtype is deprecated; rather than dtype=arr use dtype=arr.dtype. In the future this will result in an error.
    key = np.dtype(dtype)

Done pyro-ppl/numpyro#2146