Add MNIST example with SPMD for JAX #1847

Job	Run time
integration-test (v1.30.6, none, 3.10)	30m 50s
integration-test (v1.30.6, none, 3.10)	30m 43s
integration-test (v1.29.2, none, 3.10)	30m 49s
integration-test (v1.29.2, none, 3.10)	30m 44s
integration-test (v1.28.7, none, 3.8)	31m 11s
integration-test (v1.28.7, none, 3.8)	30m 50s
integration-test (v1.30.6, scheduler-plugins, 3.10)	46m 58s
integration-test (v1.30.6, scheduler-plugins, 3.10)	46m 58s
integration-test (v1.29.2, scheduler-plugins, 3.9)	47m 47s
integration-test (v1.29.2, scheduler-plugins, 3.9)	47m 47s
integration-test (v1.28.7, scheduler-plugins, 3.11)	47m 48s
integration-test (v1.28.7, scheduler-plugins, 3.11)	47m 48s
integration-test (v1.30.6, volcano, 3.10)	48m 17s
integration-test (v1.30.6, volcano, 3.10)	48m 17s
integration-test (v1.29.2, volcano, 3.8)	46m 42s
integration-test (v1.29.2, volcano, 3.8)	46m 42s
integration-test (v1.28.7, volcano, 3.10)	1h 55m 39s
integration-test (v1.28.7, volcano, 3.10)	1h 55m 54s
	14h 51m 44s

Provide feedback