Fix Log and Equal is_equivalent ignoring primitive state (#4266)

Co-authored-by: Cheng <git@zcbenz.com>

Author: kapellirohith

mlx/primitives.cpp

@@ -1941,6 +1941,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,
@@ -2795,6 +2800,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,9 +975,9 @@ 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_) {
       return "NaNEqual";
@@ -1325,9 +1325,9 @@ 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

@@ -1623,6 +1623,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()