build

Add flags for PyTorch & DCP support in the Orbax checkpoint benchmark… #42837

Job	Run time
build-export (Python 3.12, jax=newest)	2m 28s
build-orbax-model (Python 3.10, jax=newest)	58s
multiprocess-checkpoint-benchmarks (Python 3.10, jax=0.6.0)	15m 7s
build-export (Python 3.11, jax=newest)	2m 23s
build-checkpoint (Python 3.10, jax=0.6.0)	4m 8s
build-orbax-model (Python 3.12, jax=newest)	57s
build-export (Python 3.10, jax=0.4.34)	2m 20s
build-orbax-model (Python 3.10, jax=0.5.0)	53s
build-checkpoint (Python 3.11, jax=newest)	4m 23s
multiprocess-unit-tests (Python 3.11, jax=newest)	1s
build-orbax-model (Python 3.11, jax=newest)	52s
build-checkpoint (Python 3.12, jax=newest)	4m 13s
build-checkpoint (Python 3.10, jax=newest)	4m 9s
build-export (Python 3.10, jax=newest)	2m 21s
	45m 13s

Provide feedback