diff --git a/jax/_src/cloud_tpu_init.py b/jax/_src/cloud_tpu_init.py
index a61f44847..5b39994c7 100644
--- a/jax/_src/cloud_tpu_init.py
+++ b/jax/_src/cloud_tpu_init.py
@@ -77,7 +77,6 @@ def cloud_tpu_init() -> None:
   running_in_cloud_tpu_vm = True
 
   os.environ.setdefault('GRPC_VERBOSITY', 'ERROR')
-  os.environ.setdefault('JAX_PLATFORMS', 'tpu,cpu')
   os.environ['TPU_ML_PLATFORM'] = 'JAX'
   os.environ['TPU_ML_PLATFORM_VERSION'] = version.__version__
   os.environ.setdefault('ENABLE_RUNTIME_UPTIME_TELEMETRY', '1')
@@ -88,6 +87,12 @@ def cloud_tpu_init() -> None:
   os.environ.setdefault('TENSORSTORE_CURL_LOW_SPEED_TIME_SECONDS', '60')
   os.environ.setdefault('TENSORSTORE_CURL_LOW_SPEED_LIMIT_BYTES', '256')
 
+  # If the JAX_PLATFORMS env variable isn't set, config.jax_platforms defaults
+  # to None. In this case, we set it to 'tpu,cpu' to ensure that JAX uses the
+  # TPU backend.
+  if config.jax_platforms.value is None:
+    config.update('jax_platforms', 'tpu,cpu')
+
   if config.jax_pjrt_client_create_options.value is None:
     config.update(
       'jax_pjrt_client_create_options',
diff --git a/tests/config_test.py b/tests/config_test.py
index 0f49d988a..5a6da3670 100644
--- a/tests/config_test.py
+++ b/tests/config_test.py
@@ -13,10 +13,10 @@
 # limitations under the License.
 
 from absl.testing import absltest
-
 import jax
-from jax._src import test_util as jtu
 from jax._src import config
+from jax._src import test_util as jtu
+from jax._src.cloud_tpu_init import cloud_tpu_init
 
 jax.config.parse_flags_with_absl()
 
@@ -25,7 +25,8 @@ jax_test_enum_config = config.enum_state(
     name='jax_test_enum_config',
     enum_values=['default', 'xxx', 'yyy'],
     default='default',
-    help='Configuration only used for tests.')
+    help='Configuration only used for tests.',
+)
 
 
 class ConfigTest(jtu.JaxTestCase):
@@ -68,6 +69,19 @@ class ConfigTest(jtu.JaxTestCase):
         pass
     self.assertEqual(jax_test_enum_config.value, 'default')
 
+  def test_cloud_tpu_init(self):
+    if not jtu.is_cloud_tpu():
+      self.skipTest('Not running on a Cloud TPU VM.')
+
+    # Context manager resets the jax_platforms config to its original value.
+    with jtu.global_config_context(jax_platforms=None):
+      cloud_tpu_init()
+      self.assertEqual(config.jax_platforms.value, 'tpu,cpu')
+
+    with jtu.global_config_context(jax_platforms='platform_A'):
+      cloud_tpu_init()
+      self.assertEqual(config.jax_platforms.value, 'platform_A')
+
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())