Some refactoring of methods to enable additional customization of BaseTrainer, including how input data to models are sharded.#345

copybara-service[bot] wants to merge 1 commit into

maingoogle-deepmind/kfac-jax:mainfrom

test_766568928google-deepmind/kfac-jax:test_766568928Copy head branch name to clipboard