Bump to MLX 0.32.0

CMakeLists.txt

@@ -35,7 +35,7 @@ else()
   FetchContent_Declare(
     mlx
     GIT_REPOSITORY "https://github.com/ml-explore/mlx.git"
-    GIT_TAG v0.31.2)
+    GIT_TAG v0.32.0)
   FetchContent_MakeAvailable(mlx)
 endif()
 

examples/CMakeLists.txt

@@ -26,3 +26,6 @@ target_link_libraries(example-gguf PUBLIC mlxc)
 
 add_executable(example-graph ${CMAKE_CURRENT_LIST_DIR}/example-graph.c)
 target_link_libraries(example-graph PUBLIC mlxc)
+
+add_executable(example-ops ${CMAKE_CURRENT_LIST_DIR}/example-ops.c)
+target_link_libraries(example-ops PUBLIC mlxc)

examples/example-ops.c

@@ -0,0 +1,77 @@
+/* Copyright © 2023-2024 Apple Inc. */
+
+#include <stdio.h>
+#include "mlx/c/mlx.h"
+
+void print_array(const char* msg, mlx_array arr) {
+  mlx_string str = mlx_string_new();
+  mlx_array_tostring(&str, arr);
+  printf("%s\n%s\n", msg, mlx_string_data(str));
+  mlx_string_free(str);
+}
+
+/* Exercises a range of array, elementwise, and linalg ops. */
+int main(void) {
+  mlx_stream stream = mlx_default_cpu_stream_new();
+
+  float data[] = {1, -2, 0, 4, -5, 6};
+  int shape[] = {2, 3};
+  mlx_array arr = mlx_array_new_data(data, shape, 2, MLX_FLOAT32);
+  mlx_array res = mlx_array_new();
+
+  mlx_positive(&res, arr, stream);
+  print_array("positive", res);
+
+  mlx_trunc(&res, arr, stream);
+  print_array("trunc", res);
+
+  mlx_count_nonzero(&res, arr, 1, false, stream);
+  print_array("count_nonzero (axis=1)", res);
+
+  mlx_diff(&res, arr, 1, 1, stream);
+  print_array("diff (n=1, axis=1)", res);
+
+  int axes[] = {1};
+  mlx_flip(&res, arr, axes, 1, stream);
+  print_array("flip (axis=1)", res);
+
+  float rhs_data[] = {6, 5, 4, 3, 2, 1};
+  mlx_array rhs = mlx_array_new_data(rhs_data, shape, 2, MLX_FLOAT32);
+  mlx_vecdot(&res, arr, rhs, 1, stream);
+  print_array("vecdot (axis=1)", res);
+
+  mlx_array lbool = mlx_array_new_bool(true);
+  mlx_array rbool = mlx_array_new_bool(false);
+  mlx_logical_xor(&res, lbool, rbool, stream);
+  print_array("logical_xor", res);
+
+  mlx_vector_array parts = mlx_vector_array_new();
+  mlx_unstack(&parts, arr, 0, stream);
+  printf("unstack (axis=0): %zu arrays\n", mlx_vector_array_size(parts));
+
+  float square_data[] = {4, 3, 6, 3};
+  int square_shape[] = {2, 2};
+  mlx_array square =
+      mlx_array_new_data(square_data, square_shape, 2, MLX_FLOAT32);
+  mlx_linalg_det(&res, square, stream);
+  print_array("linalg_det", res);
+
+  mlx_array sign = mlx_array_new();
+  mlx_array logabsdet = mlx_array_new();
+  mlx_linalg_slogdet(&sign, &logabsdet, square, stream);
+  print_array("slogdet sign", sign);
+  print_array("slogdet logabsdet", logabsdet);
+
+  mlx_array_free(arr);
+  mlx_array_free(rhs);
+  mlx_array_free(res);
+  mlx_array_free(lbool);
+  mlx_array_free(rbool);
+  mlx_array_free(square);
+  mlx_array_free(sign);
+  mlx_array_free(logabsdet);
+  mlx_vector_array_free(parts);
+  mlx_stream_free(stream);
+
+  return 0;
+}

examples/example.c

@@ -58,6 +58,11 @@ void gpu_info(void) {
   mlx_device_info_free(info);
   mlx_device_free(dev);
 
+  const char* metallib_path = NULL;
+  if (mlx_metal_get_metallib_path(&metallib_path) == 0)
+    printf(
+        "Metal library path: %s\n", metallib_path ? metallib_path : "(null)");
+
   printf("==================================================\n");
 }
 

mlx/c/linalg.cpp

@@ -57,6 +57,17 @@ extern "C" int mlx_linalg_cross(
   }
   return 0;
 }
+extern "C" int
+mlx_linalg_det(mlx_array* res, const mlx_array a, const mlx_stream s) {
+  try {
+    mlx_array_set_(
+        *res, mlx::core::linalg::det(mlx_array_get_(a), mlx_stream_get_(s)));
+  } catch (std::exception& e) {
+    mlx_error(e.what());
+    return 1;
+  }
+  return 0;
+}
 extern "C" int mlx_linalg_eig(
     mlx_array* res_0,
     mlx_array* res_1,
@@ -261,6 +272,24 @@ extern "C" int mlx_linalg_qr(
   }
   return 0;
 }
+extern "C" int mlx_linalg_slogdet(
+    mlx_array* res_0,
+    mlx_array* res_1,
+    const mlx_array a,
+    const mlx_stream s) {
+  try {
+    {
+      auto [tpl_0, tpl_1] =
+          mlx::core::linalg::slogdet(mlx_array_get_(a), mlx_stream_get_(s));
+      mlx_array_set_(*res_0, tpl_0);
+      mlx_array_set_(*res_1, tpl_1);
+    };
+  } catch (std::exception& e) {
+    mlx_error(e.what());
+    return 1;
+  }
+  return 0;
+}
 extern "C" int mlx_linalg_solve(
     mlx_array* res,
     const mlx_array a,

mlx/c/linalg.h

@@ -44,6 +44,7 @@ int mlx_linalg_cross(
     const mlx_array b,
     int axis,
     const mlx_stream s);
+int mlx_linalg_det(mlx_array* res, const mlx_array a, const mlx_stream s);
 int mlx_linalg_eig(
     mlx_array* res_0,
     mlx_array* res_1,
@@ -97,6 +98,11 @@ int mlx_linalg_qr(
     mlx_array* res_1,
     const mlx_array a,
     const mlx_stream s);
+int mlx_linalg_slogdet(
+    mlx_array* res_0,
+    mlx_array* res_1,
+    const mlx_array a,
+    const mlx_stream s);
 int mlx_linalg_solve(
     mlx_array* res,
     const mlx_array a,

mlx/c/metal.cpp

@@ -8,6 +8,15 @@
 #include "mlx/c/error.h"
 #include "mlx/c/private/mlx.h"
 
+extern "C" int mlx_metal_get_metallib_path(const char** res) {
+  try {
+    *res = mlx::core::metal::get_metallib_path().c_str();
+  } catch (std::exception& e) {
+    mlx_error(e.what());
+    return 1;
+  }
+  return 0;
+}
 extern "C" int mlx_metal_is_available(bool* res) {
   try {
     *res = mlx::core::metal::is_available();
@@ -17,6 +26,15 @@ extern "C" int mlx_metal_is_available(bool* res) {
   }
   return 0;
 }
+extern "C" int mlx_metal_set_metallib_path(const char* path) {
+  try {
+    mlx::core::metal::set_metallib_path(std::string(path));
+  } catch (std::exception& e) {
+    mlx_error(e.what());
+    return 1;
+  }
+  return 0;
+}
 extern "C" int mlx_metal_start_capture(const char* path) {
   try {
     mlx::core::metal::start_capture(std::string(path));

mlx/c/metal.h

@@ -28,7 +28,9 @@ extern "C" {
  */
 /**@{*/
 
+int mlx_metal_get_metallib_path(const char** res);
 int mlx_metal_is_available(bool* res);
+int mlx_metal_set_metallib_path(const char* path);
 int mlx_metal_start_capture(const char* path);
 int mlx_metal_stop_capture(void);