[do not merge] variants tests #243

Job	Run time
build_and_push_jax_docker_image	18m 58s
build_and_push_pytorch_docker_image	16m 45s
imagenet_resnet_jax	3m 16s
fastmri_jax	5m 55s
imagenet_vit_jax	3m 21s
criteo_jax	25m 17s
librispeech_conformer_jax	7m 12s
librispeech_deepspeech_jax	9m 16s
wmt_jax	31m 11s
ogbg_jax	8m 29s
fastmri_pytorch	7m 52s
imagenet_resnet_pytorch	8m 25s
imagenet_vit_pytorch	15m 50s
librispeech_deepspeech_pytorch	5m 24s
wmt_pytorch	30m 15s
criteo_pytorch	25m 36s
librispeech_conformer_pytorch	2m 18s
ogbg_pytorch	5m 9s
	3h 50m 29s

Provide feedback