From 89ac0c76938a91d5ebe446e24bd47037138ad5ed Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Mon, 8 May 2023 14:38:09 -0700
Subject: [PATCH] First pass at making error_if work with sharding. Waiting on
 JAX bugs to be fixed before taking this further. TODO: make sure that things
 don't get replicated.

---
 equinox/internal/_errors.py | 33 +++++++++++++++++++++++++++++++--
 tests/conftest.py           |  2 ++
 tests/test_errors.py        | 16 ++++++++++++++++
 tests/test_pmap.py          | 11 ++++-------
 4 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/equinox/internal/_errors.py b/equinox/internal/_errors.py
index d08d28b4..73f14c90 100644
--- a/equinox/internal/_errors.py
+++ b/equinox/internal/_errors.py
@@ -2,6 +2,7 @@
 
 import jax
 import jax.core
+import jax.experimental.custom_partitioning as custom_partitioning
 import jax.interpreters.ad as ad
 import jax.interpreters.batching as batching
 import jax.interpreters.mlir as mlir
@@ -13,12 +14,40 @@
 from ._unvmap import unvmap_any, unvmap_max
 
 
-def _error_impl(pred, index, *x, msgs):
+def _raises_impl(msgs, index, x):
     def raises(_index):
         raise RuntimeError(msgs[_index.item()])
 
     struct = jax.eval_shape(lambda: x)
-    return lax.cond(pred, lambda: jax.pure_callback(raises, struct, index), lambda: x)
+    return jax.pure_callback(raises, struct, index)
+
+
+def _partition_raises(msgs, arg_shapes, arg_shardings, result_shape, result_sharding):
+    return _raises_impl, result_sharding, arg_shardings
+
+
+def _infer_sharding_from_operands_raises(msgs, arg_shapes, arg_shardings, shape):
+    del msgs, arg_shapes, shape
+    _, output_sharding = arg_shardings
+    return output_sharding
+
+
+_raises = custom_partitioning.custom_partitioning(_raises_impl, static_argnums=(0,))
+_raises.def_partition(_partition_raises, _infer_sharding_from_operands_raises)
+
+
+def _error_impl(pred, index, *x, msgs):
+    def _raises_wrapper(x):
+        # Iterate to avoid a JAX bug that means `custom_partitioning` has to return a
+        # single output.
+        # TODO: just use ft.partial(_raises, msgs, index) once this is no logner
+        # necessary.
+        new_x = []
+        for xi in x:
+            new_x.append(_raises(msgs, index, xi))
+        return tuple(new_x)
+
+    return lax.cond(pred, _raises_wrapper, lambda x: x, x)
 
 
 def _error_abstract(pred, index, *x, msgs):
diff --git a/tests/conftest.py b/tests/conftest.py
index ab3392bf..f99d12e3 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,3 +1,4 @@
+import os
 import random
 import typing
 import warnings
@@ -12,6 +13,7 @@
     "ignore",
     category=beartype.roar.BeartypeDecorHintPep585DeprecationWarning,  # pyright: ignore
 )
+os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=2"
 
 
 @pytest.fixture()
diff --git a/tests/test_errors.py b/tests/test_errors.py
index 73f007e7..090d70ea 100644
--- a/tests/test_errors.py
+++ b/tests/test_errors.py
@@ -1,5 +1,7 @@
 import jax
+import jax.experimental.mesh_utils as mesh_utils
 import jax.numpy as jnp
+import jax.sharding as sharding
 import pytest
 
 import equinox.internal as eqxi
@@ -77,3 +79,17 @@ def f(x, y, z):
         return y
 
     f(1.0, 1.0, True)
+
+
+def test_sharding():
+    x = jnp.array([0, 1])
+    mesh = sharding.Mesh(mesh_utils.create_device_mesh((2,)), ["i"])
+    spec = sharding.PartitionSpec("i")
+    named_sharding = sharding.NamedSharding(mesh, spec)
+    x = jax.device_put(x, named_sharding)
+
+    @jax.jit
+    def f(x):
+        return eqxi.error_if(x, x > 5, "oh no")
+
+    f(x)
diff --git a/tests/test_pmap.py b/tests/test_pmap.py
index 632de2bc..26dbd2b7 100644
--- a/tests/test_pmap.py
+++ b/tests/test_pmap.py
@@ -12,8 +12,8 @@
 from .helpers import shaped_allclose as _shaped_allclose
 
 
-(cpu,) = jax.devices("cpu")
-filter_pmap: Any = ft.partial(eqx.filter_pmap, devices=[cpu])
+devices = (jax.devices("cpu")[0],)
+filter_pmap: Any = ft.partial(eqx.filter_pmap, devices=devices)
 
 
 def shaped_allclose(x, y, **kwargs):
@@ -203,15 +203,12 @@ def f(x):
         y = x + 1
         return jax.lax.psum(y, axis_name="device")
 
-    n = jax.local_device_count()
-    output = filter_pmap(f, axis_name="device")(jnp.zeros(n))
+    output = filter_pmap(f, axis_name="device")(jnp.zeros(1))
 
-    assert shaped_allclose(output, n * jnp.ones(n))
+    assert shaped_allclose(output, jnp.ones(1))
 
 
 def test_map_non_jax():
-    devices = jax.local_devices()
-
     # this contains a non-jax value for the `activation` field
     # and will therefore break filter_pmap if not filtered out
     # at input and output