GPU discovery (#3055)

Co-authored-by: Angelos Katharopoulos <[email protected]>

python/src/device.cpp

@@ -1,9 +1,11 @@
-// Copyright © 2023-2024 Apple Inc.
+// Copyright © 2023-2025 Apple Inc.
 
 #include <sstream>
 
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/string.h>
+#include <nanobind/stl/unordered_map.h>
+#include <nanobind/stl/variant.h>
 
 #include "mlx/device.h"
 #include "mlx/utils.h"
@@ -63,4 +65,34 @@ void init_device(nb::module_& m) {
       &mx::is_available,
       "device"_a,
       R"pbdoc(Check if a back-end is available for the given device.)pbdoc");
+  m.def(
+      "device_count",
+      &mx::device_count,
+      "device_type"_a,
+      R"pbdoc(
+      Get the number of available devices for the given device type.
+
+      Args:
+          device_type (DeviceType): The type of device to query (cpu or gpu).
+
+      Returns:
+          int: Number of devices.
+      )pbdoc");
+  m.def(
+      "device_info",
+      &mx::device_info,
+      nb::arg("d") = mx::default_device(),
+      R"pbdoc(
+      Get information about a device.
+
+      Returns a dictionary with device properties. Available keys depend
+      on the backend and device type. Common keys include ``device_name``,
+      ``architecture``, and ``total_memory`` (or ``memory_size``).
+
+      Args:
+          d (Device): The device to query (defaults to the default device).
+
+      Returns:
+          dict: Device information.
+      )pbdoc");
 }

python/src/metal.cpp

@@ -9,6 +9,7 @@
 #include <nanobind/stl/vector.h>
 
 #include "mlx/backend/metal/metal.h"
+#include "mlx/device.h"
 #include "mlx/memory.h"
 #include "python/src/small_vector.h"
 
@@ -90,21 +91,8 @@ void init_metal(nb::module_& m) {
       R"pbdoc(
       Stop a Metal capture.
       )pbdoc");
-  metal.def(
-      "device_info",
-      &mx::metal::device_info,
-      R"pbdoc(
-      Get information about the GPU device and system settings.
-
-      Currently returns:
-
-      * ``architecture``
-      * ``max_buffer_size``
-      * ``max_recommended_working_set_size``
-      * ``memory_size``
-      * ``resource_limit``
-
-      Returns:
-          dict: A dictionary with string keys and string or integer values.
-      )pbdoc");
+  metal.def("device_info", []() {
+    DEPRECATE("mx.metal.device_info", "mx.device_info");
+    return mx::device_info(mx::Device(mx::Device::gpu, 0));
+  });
 }