Refactor(JAX): Use jax.jit's out_shardings instead of _enforce_jax_state_sharding#21559

Merged

hertschuh merged 14 commits intokeras-team:masterfrom

buildwithsuhana:jax_out_sharding

Aug 14, 2025

Commits on Aug 4, 2025

Refactor JAXTrainer sharding to use out_shardings
buildwithsuhana
committed

Commits on Aug 6, 2025

Commits on Aug 7, 2025

Merge branch 'keras-team:master' into jax_out_sharding
buildwithsuhana
authored

Commits on Aug 8, 2025

Update keras/src/trainers/trainer_test.py for nit changes

buildwithsuhana
and
gemini-code-assist[bot]
authored
Added a helper to reduce code duplication
buildwithsuhana
committed

Commits on Aug 9, 2025

Updated the test and the out sharding logic
buildwithsuhana
committed

Commits on Aug 11, 2025

Removing clear jax state sharding
buildwithsuhana
committed

Commits on Aug 13, 2025