Removing references to jax.config.spmd_mode('allow_all').

[pschuh]

This flag no longer does anything in jax.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 82.68%. Comparing base (2111fbc) to head (496a2f1).

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21164      +/-   ##
==========================================
- Coverage   82.69%   82.68%   -0.01%     
==========================================
  Files         564      564              
  Lines       54223    54221       -2     
  Branches     8424     8424              
==========================================
- Hits        44837    44835       -2     
  Misses       7310     7310              
  Partials     2076     2076              

Flag	Coverage Δ
keras	82.49% <100.00%> (-0.01%)	⬇️
keras-jax	63.92% <100.00%> (-0.01%)	⬇️
keras-numpy	59.03% <0.00%> (+<0.01%)	⬆️
keras-openvino	32.98% <0.00%> (+<0.01%)	⬆️
keras-tensorflow	64.30% <0.00%> (+<0.01%)	⬆️
keras-torch	63.99% <0.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.