diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 3c3d4fc604..b357b05862 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -1939,6 +1939,11 @@ std::pair, std::vector> Equal::vmap( return {{equal(a, b, stream())}, {to_ax}}; } +bool Equal::is_equivalent(const Primitive& other) const { + const Equal& e_other = static_cast(other); + return equal_nan_ == e_other.equal_nan_; +} + std::vector Equal::vjp( const std::vector& primals, const std::vector& cotangents, @@ -2791,6 +2796,11 @@ std::pair, std::vector> Log::vmap( axes}; } +bool Log::is_equivalent(const Primitive& other) const { + const Log& l_other = static_cast(other); + return base_ == l_other.base_; +} + std::vector Log1p::vjp( const std::vector& primals, const std::vector& cotangents, diff --git a/mlx/primitives.h b/mlx/primitives.h index 3a3d0ba5e5..7d2d2bccaa 100644 --- a/mlx/primitives.h +++ b/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_; diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index 76e8916538..0b1e9bbbe4 100644 --- a/python/tests/test_compile.py +++ b/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() diff --git a/tests/compile_tests.cpp b/tests/compile_tests.cpp index 30c2f887ac..14f7f1f9ef 100644 --- a/tests/compile_tests.cpp +++ b/tests/compile_tests.cpp @@ -212,6 +212,37 @@ TEST_CASE("test no simplify") { set_compile_mode(CompileMode::enabled); } +auto log_bases(const std::vector& inputs) { + auto a = inputs[0]; + return std::vector{log(a) + log2(a)}; +}; + +auto equal_nan_variants(const std::vector& inputs) { + auto a = inputs[0]; + return std::vector{ + 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()); + + auto c = array({1.0f, std::numeric_limits::quiet_NaN()}); + auto d = compile(equal_nan_variants)({c})[0]; + CHECK(array_equal(d, array({false, true})).item()); + + // Matching state still simplifies. + auto same_base = [](const std::vector& inputs) -> std::vector { + 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&) { auto a = array(1.0); auto b = array(2.0);