Fix Log and Equal is_equivalent ignoring primitive state

mlx/primitives.cpp

@@ -1939,6 +1939,11 @@ std::pair<std::vector<array>, std::vector<int>> Equal::vmap(
   return {{equal(a, b, stream())}, {to_ax}};
 }
 
+bool Equal::is_equivalent(const Primitive& other) const {
+  const Equal& e_other = static_cast<const Equal&>(other);
+  return equal_nan_ == e_other.equal_nan_;
+}
+
 std::vector<array> Equal::vjp(
     const std::vector<array>& primals,
     const std::vector<array>& cotangents,
@@ -2791,6 +2796,11 @@ std::pair<std::vector<array>, std::vector<int>> Log::vmap(
       axes};
 }
 
+bool Log::is_equivalent(const Primitive& other) const {
+  const Log& l_other = static_cast<const Log&>(other);
+  return base_ == l_other.base_;
+}
+
 std::vector<array> Log1p::vjp(
     const std::vector<array>& primals,
     const std::vector<array>& cotangents,

mlx/primitives.h

@@ -975,8 +975,8 @@ class Equal : public UnaryPrimitive {
 
   DEFINE_VMAP()
   DEFINE_GRADS()
-  DEFINE_DEFAULT_IS_EQUIVALENT()
   DEFINE_INPUT_OUTPUT_SHAPE()
+  bool is_equivalent(const Primitive& other) const override;
 
   const char* name() const override {
     if (equal_nan_) {
@@ -1325,8 +1325,8 @@ class Log : public UnaryPrimitive {
 
   DEFINE_VMAP()
   DEFINE_GRADS()
-  DEFINE_DEFAULT_IS_EQUIVALENT()
   DEFINE_INPUT_OUTPUT_SHAPE()
+  bool is_equivalent(const Primitive& other) const override;
 
   Base state() const {
     return base_;

python/tests/test_compile.py

@@ -1576,6 +1576,29 @@ def test_compile_abs_unsigned(self):
             x = mx.array([1, 2, 3], dtype)
             self.assertTrue(mx.array_equal(mx.compile(fun)(x), fun(x)))
 
+    def test_compile_different_log_bases(self):
+        # The logs are intermediates, since outputs are not simplified.
+        def entropies(p):
+            nats = -mx.sum(p * mx.log(p))
+            bits = -mx.sum(p * mx.log2(p))
+            return mx.stack([nats, bits])
+
+        p = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32)
+        expected = np.array(
+            [-(p * np.log(p)).sum(), -(p * np.log2(p)).sum()], dtype=np.float32
+        )
+        out = mx.compile(entropies)(mx.array(p))
+        self.assertTrue(np.allclose(out, expected, atol=1e-5))
+
+    def test_compile_equal_nan(self):
+        def fun(x):
+            return mx.stack(
+                [mx.array_equal(x, x), mx.array_equal(x, x, equal_nan=True)]
+            )
+
+        x = mx.array([1.0, float("nan"), 3.0])
+        self.assertTrue(mx.array_equal(mx.compile(fun)(x), mx.array([False, True])))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()

tests/compile_tests.cpp

@@ -212,6 +212,37 @@ TEST_CASE("test no simplify") {
   set_compile_mode(CompileMode::enabled);
 }
 
+auto log_bases(const std::vector<array>& inputs) {
+  auto a = inputs[0];
+  return std::vector<array>{log(a) + log2(a)};
+};
+
+auto equal_nan_variants(const std::vector<array>& inputs) {
+  auto a = inputs[0];
+  return std::vector<array>{
+      stack({array_equal(a, a), array_equal(a, a, true)})};
+};
+
+TEST_CASE("test no simplify different primitive state") {
+  set_compile_mode(CompileMode::no_fuse);
+  auto a = array({2.0f, 8.0f});
+  auto b = compile(log_bases)({a})[0];
+  CHECK(b.inputs()[0].id() != b.inputs()[1].id());
+  CHECK(allclose(b, log(a) + log2(a)).item<bool>());
+
+  auto c = array({1.0f, std::numeric_limits<float>::quiet_NaN()});
+  auto d = compile(equal_nan_variants)({c})[0];
+  CHECK(array_equal(d, array({false, true})).item<bool>());
+
+  // Matching state still simplifies.
+  auto same_base = [](const std::vector<array>& inputs) -> std::vector<array> {
+    return {log(inputs[0]) + log(inputs[0])};
+  };
+  auto e = compile(same_base)({a})[0];
+  CHECK(e.inputs()[0].id() == e.inputs()[1].id());
+  set_compile_mode(CompileMode::enabled);
+}
+
 auto multi_one(const std::vector<array>&) {
   auto a = array(1.0);
   auto b = array(2.0);