Repository: jax-ml/jax
Pull request #24992 (2024-11-20 01:51:42)

Fix a bug where constant deduplication used an inappropriate inequality.
Fix a bug where constant deduplication used an inappropriate inequality.

We need to compare constants for bitwise equality, not, e.g., floating point equality. The change that added deduplication caused us to conflate +0.0 and -0.0, which led a downstream test not to terminate.



Patch:
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index 23325a9d7..102e4f490 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -1755,33 +1755,39 @@ def _emit_lowering_rule_as_fun(lowering_rule,
 class HashableLiteral:
   """Hashable wrapper of core.Literal, used for deduplicating IR constants."""
 
-  __slots__ = ["value"]
+  __slots__ = ["value", "data"]
 
   value: core.Literal
 
+  # Copy of the value suitable for an equality comparison. We are careful to
+  # avoid floating point comparisons here, because in particular we don't want
+  # 0.0 and -0.0 to be considered equal, but we are fine with NaNs being equal.
+  data: bytes | int | bool | None
+
   def __init__(self, value):
     self.value = value
+    if isinstance(value.val, (np.generic, np.ndarray)):
+      self.data = value.val.tobytes()
+    elif isinstance(value.val, (bool, int)):
+      self.data = value.val
+    elif isinstance(value.val, float):
+      self.data = np.float64(value.val).tobytes()
+    elif isinstance(value.val, complex):
+      self.data = np.complex128(value.val).tobytes()
+    else:
+      self.data = None  # Unhandled case.
 
   def __hash__(self):
-    h = self.value.hash
-    return id(self.value.val) if h is None else h
+    return hash(self.data)
 
   def __eq__(self, other):
-    if self is other:
-      return True
     if type(self.value.val) != type(other.value.val):
       return False
     if self.value.aval != other.value.aval:
       return False
-    if isinstance(self.value.val, (bool, int, float, complex)):
-      return self.value == other.value
-    if isinstance(self.value.val, (np.generic, np.ndarray)):
-      return np.array_equal(
-        self.value.val, other.value.val,
-        equal_nan=np.issubdtype(self.value.val.dtype, np.inexact))
-    # Since the use case is constant deduplication, it's safe to return
-    # False in unhandled cases.
-    return False
+    if self.data is None:
+      return id(self) == id(other)
+    return self.data == other.data
 
 
 def jaxpr_subcomp(ctx: ModuleContext, jaxpr: core.Jaxpr,


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 6df011419..e32424cfd 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -1312,6 +1312,17 @@ class PJitTest(jtu.BufferDonationTestCase):
     ans2 = jnp.sin(x0) # cpp_pjit hit with bad cache entry
     assert(ans1.devices() == ans2.devices())
 
+  def test_zero_literal_equality(self):
+    # This test verifies that we don't accidentally conflate positive and
+    # negative zeros when deduplicating literals in the IR.
+    f = jax.jit(lambda x: (x / np.float32(-0.0), x / np.float32(0.0)))
+    a, b = f(np.float32(1.0))
+    self.assertEqual(a, -np.inf)
+    self.assertEqual(b, np.inf)
+    ir = f.lower(np.float32(1.0)).as_text()
+    self.assertIn("stablehlo.constant dense<0.000000e+00>", ir)
+    self.assertIn("stablehlo.constant dense<-0.000000e+00>", ir)
+
 @jtu.pytest_mark_if_available('multiaccelerator')
 class CustomPartitionerTest(jtu.JaxTestCase):