Fix compile hasher for string constants. (#1677)

* fix hash * add test * nit
ml-explore · Dec 9, 2024 · 35b412c · 35b412c
1 parent d0f471c
commit 35b412c
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 6 deletions.
diff --git a/python/src/transforms.cpp b/python/src/transforms.cpp
@@ -433,19 +433,18 @@ struct PyCompiledFun {
         auto d = nb::cast<nb::dict>(obj);
         constants.push_back(dict_identifier);
         for (auto item : d) {
-          auto r = item.first.attr("__hash__");
-          constants.push_back(*reinterpret_cast<uint64_t*>(&r));
+          auto r = item.first.attr("__hash__")();
+          constants.push_back(nb::cast<int64_t>(r));
           recurse(item.second);
         }
       } else if (nb::isinstance<array>(obj)) {
         inputs.push_back(nb::cast<array>(obj));
         constants.push_back(array_identifier);
       } else if (nb::isinstance<nb::str>(obj)) {
-        auto r = obj.attr("__hash__");
-        constants.push_back(*reinterpret_cast<uint64_t*>(&r));
+        auto r = obj.attr("__hash__")();
+        constants.push_back(nb::cast<int64_t>(r));
       } else if (nb::isinstance<nb::int_>(obj)) {
-        auto r = nb::cast<int64_t>(obj);
-        constants.push_back(*reinterpret_cast<uint64_t*>(&r));
+        constants.push_back(nb::cast<int64_t>(obj));
       } else if (nb::isinstance<nb::float_>(obj)) {
         auto r = nb::cast<double>(obj);
         constants.push_back(*reinterpret_cast<uint64_t*>(&r));

diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py
@@ -579,6 +579,27 @@ def fun(x, y):
 
         self.assertEqual(counter[0], 2)
 
+        y = 1.0
+
+        @mx.compile
+        def fun(x, constant):
+            return x + y
+
+        constant1 = "abc"
+        out = fun(mx.array(0.0), constant1)
+        self.assertEqual(out, mx.array(1.0))
+
+        # new object, same value, no recompilation
+        y = 2.0
+        constant2 = "abc".encode("utf-8").decode("utf-8")
+        out = fun(mx.array(0.0), constant2)
+        self.assertEqual(out, mx.array(1.0))
+
+        # same object, new value, recompilation
+        constant2 = "xyz"
+        out = fun(mx.array(0.0), constant2)
+        self.assertEqual(out, mx.array(2.0))
+
     def test_compile_inf(self):
         @mx.compile
         def fun(x):