Upgrade Jax environment and fix workload compatibility#1158

alfredyu-cienet merged 2 commits intoGoogleCloudPlatform:masterGoogleCloudPlatform/ml-auto-solutions:masterfrom

CIeNET-International:tpu-obs/user/yuna/jax-workload-updateCIeNET-International/ml-auto-solutions-3:tpu-obs/user/yuna/jax-workload-updateCopy head branch name to clipboard