From 935c9859d64f30f17ee34f859ef9f09a7488aed3 Mon Sep 17 00:00:00 2001
From: Cheng <git@zcbenz.com>
Date: Thu, 19 Mar 2026 16:59:08 +0900
Subject: [PATCH] Do not create device when not using device

---
 python/src/device.cpp     | 8 ++++++--
 python/src/transforms.cpp | 5 -----
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/python/src/device.cpp b/python/src/device.cpp
index f15f7f92db..e70b69bd34 100644
--- a/python/src/device.cpp
+++ b/python/src/device.cpp
@@ -1,8 +1,10 @@
 // Copyright © 2023-2025 Apple Inc.
 
+#include <optional>
 #include <sstream>
 
 #include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/unordered_map.h>
 #include <nanobind/stl/variant.h>
@@ -80,8 +82,10 @@ void init_device(nb::module_& m) {
       )pbdoc");
   m.def(
       "device_info",
-      &mx::device_info,
-      nb::arg("d") = mx::default_device(),
+      [](std::optional<mx::Device> d) {
+        return mx::device_info(d.value_or(mx::default_device()));
+      },
+      "d"_a = nb::none(),
       R"pbdoc(
       Get information about a device.
 
diff --git a/python/src/transforms.cpp b/python/src/transforms.cpp
index c615b12ec5..2168d4bbf6 100644
--- a/python/src/transforms.cpp
+++ b/python/src/transforms.cpp
@@ -1534,9 +1534,4 @@ void init_transforms(nb::module_& m) {
           A callable that recomputes intermediate states during gradient
           computation.
       )pbdoc");
-
-  // Register static Python object cleanup before the interpreter exits
-  auto atexit = nb::module_::import_("atexit");
-  atexit.attr("register")(
-      nb::cpp_function([]() { mx::detail::compile_clear_cache(); }));
 }