Repository: jax-ml/jax
Pull request #24751 (2024-11-07 00:36:50)

Allow setting default_device with platform names.

Stella-S-Yan: Thank you both for reviewing my first PR :).  I've made the requested changes and updated the PR. Please take another look. 

Patch:
diff --git a/jax/_src/config.py b/jax/_src/config.py
index f3edde699..30a9ba0be 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -1561,7 +1561,9 @@ def _update_default_device_thread_local(val):
 
 
 def _validate_default_device(val):
-  if val is not None and not isinstance(val, xla_client.Device):
+  if (val is not None and
+      not isinstance(val, xla_client.Device) and
+      val not in ['cpu', 'gpu', 'tpu']):
     # TODO(skyewm): this is a workaround for non-PJRT Device types. Remove when
     # all JAX backends use a single C++ device interface.
     if 'Device' in str(type(val)):
@@ -1569,12 +1571,11 @@ def _validate_default_device(val):
           'Allowing non-`xla_client.Device` default device: %s, type: %s',
           repr(val), type(val))
       return
-    raise ValueError('jax.default_device must be passed a Device object (e.g. '
-                     f"`jax.devices('cpu')[0]`), got: {val!r}")
+    raise ValueError('jax.default_device must be passed either a Device object (e.g. '
+                     f"`jax.devices('cpu')[0]`) or a platform name string like 'cpu' or 'gpu'"
+                     f", got: {val!r}")
 
 
-# TODO(skye): default_device only accepts devices for now. Make it work with
-# platform names as well (e.g. "cpu" to mean the same as jax.devices("cpu")[0]).
 default_device = string_or_object_state(
     name='jax_default_device',
     default=None,
diff --git a/jax/_src/interpreters/pxla.py b/jax/_src/interpreters/pxla.py
index c83d3e3a4..2ee2ec75e 100644
--- a/jax/_src/interpreters/pxla.py
+++ b/jax/_src/interpreters/pxla.py
@@ -1710,7 +1710,10 @@ ShardingInfo = tuple[
 
 
 def _get_default_device() -> xc.Device:
-  return config.default_device.value or xb.local_devices()[0]
+  if isinstance(config.default_device.value, str):
+    return xb.get_backend(config.default_device.value).local_devices()[0]
+  else:
+    return config.default_device.value or xb.local_devices()[0]
 
 
 def _get_and_check_device_assignment(
@@ -1742,6 +1745,7 @@ def _get_and_check_device_assignment(
         raise DeviceAssignmentMismatchError([
             DeviceAssignmentMismatch(devices, MismatchType.CONTEXT_DEVICES, None),
             DeviceAssignmentMismatch(arr_device_assignment, s_type, source_info)])
+
   if first_sharding_info is None and devices:
     final_device_assignment = devices
   elif first_sharding_info is None:
@@ -2190,6 +2194,7 @@ def lower_sharding_computation(
   assert len(out_shardings) == len(out_layouts) == len(global_out_avals), (
       len(out_shardings), len(out_layouts), len(global_out_avals))
 
+
   devices_from_context = (None if context_mesh is None or context_mesh.empty
                           else context_mesh._flat_devices_tuple)
   # Device assignment across all inputs, outputs and shardings inside jaxpr
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
index f74c74077..ace06be47 100644
--- a/third_party/xla/workspace.bzl
+++ b/third_party/xla/workspace.bzl
@@ -21,8 +21,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 #    curl -L https://github.com/openxla/xla/archive/<git hash>.tar.gz | sha256sum
 #    and update XLA_SHA256 with the result.
 
-XLA_COMMIT = "e93a258e4494231626c7d3b6a6447e746ea72f9c"
-XLA_SHA256 = "99f3a6b06230becf013f00009afeee4c89f52818e7a4a1ea4851157dc853830e"
+XLA_COMMIT = "cebb50cc77915aaf16e76b10d78e4c3deb955939"
+XLA_SHA256 = "e4be11c05a6b59e8a090e6205d34f138f889826b74113633822fd11b65258668"
 
 def repo():
     tf_http_archive(


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
index 8ab5d90f6..e61236a2d 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -287,13 +287,14 @@ class JitTest(jtu.BufferDonationTestCase):
     self.assertEqual(f(sticky).devices(), system_default_devices)
     self.assertEqual(f(1).devices(), system_default_devices)
 
-  # TODO(skye): make this work!
   def test_jit_default_platform(self):
-    with self.assertRaisesWithLiteralMatch(
-        ValueError, "jax.default_device must be passed a Device object "
-        "(e.g. `jax.devices('cpu')[0]`), got: 'cpu'"):
       with jax.default_device("cpu"):
-        jax.jit(lambda x: x + 1)(1)
+        result = jax.jit(lambda x: x + 1)(1)
+      self.assertEqual(result.device.platform, "cpu")
+
+      result = jax.jit(lambda x: x + 1)(1)
+      self.assertEqual(result.device.platform, jax.default_backend())
+
 
   def test_complex_support(self):
     self.assertEqual(jit(lambda x: x + 1)(1 + 1j), 2 + 1j)