Repository: google/jax
Pull request #21978 (2024-06-19 11:52:34)

[export] Fix multi-platform lowering for unknown platform, with donated_argnums

I had to ensure that the check for platforms supporting donation only kicks in when we actually have donation.


Patch:
diff --git a/jax/_src/export/_export.py b/jax/_src/export/_export.py
index 0689f98b4..efd240aa9 100644
--- a/jax/_src/export/_export.py
+++ b/jax/_src/export/_export.py
@@ -161,10 +161,10 @@ class Exported:
         the mesh. See `out_shardings_jax` for a way to turn these
         into sharding specification that can be used with JAX APIs.
     nr_devices: the number of devices that the module has been lowered for.
-    platforms: a tuple containing at least one of 'tpu', 'cpu',
-        'cuda', 'rocm'. See https://jax.readthedocs.io/en/latest/export.html#module-calling-convention
-        for the calling convention for when
-        there are multiple export platforms.
+    platforms: a tuple containing the platforms for which the function should
+        be exported. The set of platforms in JAX is open-ended; users can
+        add platforms. JAX built-in platforms are: 'tpu', 'cpu', 'cuda', 'rocm'.
+        See https://jax.readthedocs.io/en/latest/export/export.html#cross-platform-and-multi-platform-export.
     ordered_effects: the ordered effects present in the serialized module.
         This is present from serialization version 9. See https://jax.readthedocs.io/en/latest/export.html#module-calling-convention
         for the calling convention in presence of ordered effects.
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index 3ba40907b..251888966 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -886,7 +886,8 @@ def lower_jaxpr_to_module(
   platforms_with_donation = [p for p in platforms
                              if p in _platforms_with_donation]
   if platforms_with_donation:
-    if len(platforms_with_donation) != len(platforms):
+    if len(platforms_with_donation) != len(platforms) and (
+        xla_donated_args or any(donated_args)):
       raise NotImplementedError(
         "In multi-platform lowering either all or no lowering platforms "
         f"should support donation. Lowering for {platforms} of which "
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 6b6671449..71357b4a8 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -127,37 +127,8 @@ get_printoptions = np.get_printoptions
 printoptions = np.printoptions
 set_printoptions = np.set_printoptions
 
+@util.implements(np.iscomplexobj)
 def iscomplexobj(x: Any) -> bool:
-  """Check if the input is a complex number or an array containing complex elements.
-
-  JAX implementation of :func:`numpy.iscomplexobj`.
-
-  The function evaluates based on input type rather than value.
-  Inputs with zero imaginary parts are still considered complex.
-
-  Args:
-    x: input object to check.
-
-  Returns:
-    True if ``x`` is a complex number or an array containing at least one complex element,
-    False otherwise.
-
-  See Also:
-    - :func:`jax.numpy.isrealobj`
-    - :func:`jax.numpy.iscomplex`
-
-  Examples:
-    >>> jnp.iscomplexobj(True)
-    False
-    >>> jnp.iscomplexobj(0)
-    False
-    >>> jnp.iscomplexobj(jnp.array([1, 2]))
-    False
-    >>> jnp.iscomplexobj(1+2j)
-    True
-    >>> jnp.iscomplexobj(jnp.array([0, 1+2j]))
-    True
-  """
   if x is None:
     return False
   try:


Test patch:
diff --git a/tests/export_test.py b/tests/export_test.py
index e9c58dea9..d93e07f1f 100644
--- a/tests/export_test.py
+++ b/tests/export_test.py
@@ -911,6 +911,17 @@ class JaxExportTest(jtu.JaxTestCase):
     a = exp2.in_avals[0].shape[0]
     self.assertEqual(exp2.out_avals[0].shape, output_shape(a))
 
+  def test_with_donation(self):
+    f = jax.jit(jnp.sin, donate_argnums=(0,))
+    x = np.arange(3, dtype=np.float32)
+    exp = export.export(f)(x)
+
+    def caller(x):
+      y = exp.call(x)
+      return x + y
+    res = jax.jit(caller)(x)
+    self.assertAllClose(res, x + np.sin(x))
+
   def test_poly_call_pmap(self):
     if len(jax.devices()) < 2:
       self.skipTest("Need at least 2 devices")
@@ -1312,7 +1323,7 @@ class JaxExportTest(jtu.JaxTestCase):
   def test_multi_platform(self):
     x = np.arange(8, dtype=np.float32)
     exp = get_exported(jax.jit(_testing_multi_platform_func),
-                       lowering_platforms=("tpu", "cpu", "cuda","rocm"))(x)
+                       lowering_platforms=("tpu", "cpu", "cuda", "rocm"))(x)
     self.assertEqual(exp.platforms, ("tpu", "cpu", "cuda", "rocm"))
     module_str = str(exp.mlir_module())
     expected_main_re = (
@@ -1335,14 +1346,14 @@ class JaxExportTest(jtu.JaxTestCase):
   def test_multi_platform_nested(self):
     x = np.arange(5, dtype=np.float32)
     exp = get_exported(jax.jit(lambda x: _testing_multi_platform_func(jnp.sin(x))),
-                       lowering_platforms=("cpu", "tpu", "cuda","rocm"))(x)
-    self.assertEqual(exp.platforms, ("cpu", "tpu", "cuda","rocm"))
+                       lowering_platforms=("cpu", "tpu", "cuda", "rocm"))(x)
+    self.assertEqual(exp.platforms, ("cpu", "tpu", "cuda", "rocm"))
 
     # Now serialize the call to the exported using a different sequence of
     # lowering platforms, but included in the lowering platforms for the
     # nested exported.
     exp2 = get_exported(jax.jit(exp.call),
-                        lowering_platforms=("cpu", "cuda","rocm"))(x)
+                        lowering_platforms=("cpu", "cuda", "rocm"))(x)
 
     # Ensure that we do not have multiple lowerings of the exported function
     exp2_module_str = str(exp2.mlir_module())
@@ -1361,7 +1372,7 @@ class JaxExportTest(jtu.JaxTestCase):
   def test_multi_platform_nested_inside_single_platform_export(self):
     x = np.arange(5, dtype=np.float32)
     exp = get_exported(jax.jit(_testing_multi_platform_func),
-                       lowering_platforms=("cpu", "tpu", "cuda","rocm"))(x)
+                       lowering_platforms=("cpu", "tpu", "cuda", "rocm"))(x)
     self.assertEqual(exp.platforms, ("cpu", "tpu", "cuda", "rocm"))
 
     # Now serialize the call for the current platform.
@@ -1430,6 +1441,31 @@ class JaxExportTest(jtu.JaxTestCase):
     expected = x * np.float32(dict(cpu=2, gpu=3, tpu=4)[jtu.device_under_test()])
     self.assertAllClose(exp.call(x), expected)
 
+  def test_multi_platform_unknown_platform(self):
+    x = np.arange(8, dtype=np.float32)
+    exp = get_exported(jax.jit(jnp.sin),
+                       lowering_platforms=("tpu", "cpu", "cuda", "other"))(x)
+    self.assertEqual(exp.platforms, ("tpu", "cpu", "cuda", "other"))
+
+
+  def test_multi_platform_with_donation(self):
+    f = jax.jit(jnp.sin, donate_argnums=(0,))
+    x = np.arange(3, dtype=np.float32)
+    exp = export.export(f, platforms=["cpu", "tpu"])(x)
+    if jtu.device_under_test() not in ["cpu", "tpu"]:
+      self.skipTest("other platform")
+
+    def caller(x):
+      y = exp.call(x)
+      return x + y
+    res = jax.jit(caller)(x)
+    self.assertAllClose(res, x + np.sin(x))
+
+    with self.assertRaisesRegex(
+        NotImplementedError,
+        "In multi-platform lowering either all or no lowering platforms should support donation"):
+      export.export(f, platforms=["cpu", "tpu", "other"])(x)
+
   def test_multi_platform_and_poly(self):
     if jtu.test_device_matches(["gpu"]):
       # The export is not applicable to GPU
@@ -1697,5 +1733,6 @@ class JaxExportTest(jtu.JaxTestCase):
     res_exported = exp_f.call(lhs, rhs, group_sizes)
     self.assertAllClose(res_native, res_exported)
 
+
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())