keras/src/callbacks/callback.py/Callback/model
class Callback:    def model(self):
        if backend.backend() == "jax" and hasattr(
            self._model, "jax_state_sync"
        ):
            # With JAX, by default the model state is not
            # attached to the model in the middle of an
            # epoch. We have to force a sync before
            # accessing model state for e.g. checkpointing.
            self._model.jax_state_sync()
        return self._model