Fix JAX GPU tests.

Using XLA_PYTHON_CLIENT_MEM_FRACTION.

Author: hertschuh

.kokoro/github/ubuntu/gpu/jax/continuous.cfg

@@ -12,5 +12,10 @@ env_vars: {
    value: "jax"
 }
 
+env_vars: {
+   key: "XLA_PYTHON_CLIENT_MEM_FRACTION"
+   value: ".5"
+}
+
 # Set timeout to 60 mins from default 180 mins
 timeout_mins: 60