Bump to MLX v0.26.1 (#76)

* bump to MLX v0.26.1

* variant generator: check overload/name sizes before printing

* mlxvariants: handle additional overload for normal()

* regenerate files

Author: andresy

CMakeLists.txt

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

mlx/c/fft.cpp

@@ -68,6 +68,25 @@ extern "C" int mlx_fft_fftn(
   }
   return 0;
 }
+extern "C" int mlx_fft_fftshift(
+    mlx_array* res,
+    const mlx_array a,
+    const int* axes,
+    size_t axes_num,
+    const mlx_stream s) {
+  try {
+    mlx_array_set_(
+        *res,
+        mlx::core::fft::fftshift(
+            mlx_array_get_(a),
+            std::vector<int>(axes, axes + axes_num),
+            mlx_stream_get_(s)));
+  } catch (std::exception& e) {
+    mlx_error(e.what());
+    return 1;
+  }
+  return 0;
+}
 extern "C" int mlx_fft_ifft(
     mlx_array* res,
     const mlx_array a,
@@ -128,6 +147,25 @@ extern "C" int mlx_fft_ifftn(
   }
   return 0;
 }
+extern "C" int mlx_fft_ifftshift(
+    mlx_array* res,
+    const mlx_array a,
+    const int* axes,
+    size_t axes_num,
+    const mlx_stream s) {
+  try {
+    mlx_array_set_(
+        *res,
+        mlx::core::fft::ifftshift(
+            mlx_array_get_(a),
+            std::vector<int>(axes, axes + axes_num),
+            mlx_stream_get_(s)));
+  } catch (std::exception& e) {
+    mlx_error(e.what());
+    return 1;
+  }
+  return 0;
+}
 extern "C" int mlx_fft_irfft(
     mlx_array* res,
     const mlx_array a,

mlx/c/fft.h

@@ -49,6 +49,12 @@ int mlx_fft_fftn(
     const int* axes,
     size_t axes_num,
     const mlx_stream s);
+int mlx_fft_fftshift(
+    mlx_array* res,
+    const mlx_array a,
+    const int* axes,
+    size_t axes_num,
+    const mlx_stream s);
 int mlx_fft_ifft(
     mlx_array* res,
     const mlx_array a,
@@ -71,6 +77,12 @@ int mlx_fft_ifftn(
     const int* axes,
     size_t axes_num,
     const mlx_stream s);
+int mlx_fft_ifftshift(
+    mlx_array* res,
+    const mlx_array a,
+    const int* axes,
+    size_t axes_num,
+    const mlx_stream s);
 int mlx_fft_irfft(
     mlx_array* res,
     const mlx_array a,

mlx/c/linalg.cpp

@@ -57,6 +57,24 @@ extern "C" int mlx_linalg_cross(
   }
   return 0;
 }
+extern "C" int mlx_linalg_eig(
+    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::eig(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_eigh(
     mlx_array* res_0,
     mlx_array* res_1,
@@ -76,6 +94,18 @@ extern "C" int mlx_linalg_eigh(
   }
   return 0;
 }
+extern "C" int
+mlx_linalg_eigvals(mlx_array* res, const mlx_array a, const mlx_stream s) {
+  try {
+    mlx_array_set_(
+        *res,
+        mlx::core::linalg::eigvals(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_eigvalsh(
     mlx_array* res,
     const mlx_array a,

mlx/c/linalg.h

@@ -43,12 +43,18 @@ int mlx_linalg_cross(
     const mlx_array b,
     int axis,
     const mlx_stream s);
+int mlx_linalg_eig(
+    mlx_array* res_0,
+    mlx_array* res_1,
+    const mlx_array a,
+    const mlx_stream s);
 int mlx_linalg_eigh(
     mlx_array* res_0,
     mlx_array* res_1,
     const mlx_array a,
     const char* UPLO,
     const mlx_stream s);
+int mlx_linalg_eigvals(mlx_array* res, const mlx_array a, const mlx_stream s);
 int mlx_linalg_eigvalsh(
     mlx_array* res,
     const mlx_array a,

mlx/c/ops.cpp

@@ -2752,6 +2752,26 @@ extern "C" int mlx_scatter_prod(
   }
   return 0;
 }
+extern "C" int mlx_segmented_mm(
+    mlx_array* res,
+    const mlx_array a,
+    const mlx_array b,
+    const mlx_array segments,
+    const mlx_stream s) {
+  try {
+    mlx_array_set_(
+        *res,
+        mlx::core::segmented_mm(
+            mlx_array_get_(a),
+            mlx_array_get_(b),
+            mlx_array_get_(segments),
+            mlx_stream_get_(s)));
+  } catch (std::exception& e) {
+    mlx_error(e.what());
+    return 1;
+  }
+  return 0;
+}
 extern "C" int
 mlx_sigmoid(mlx_array* res, const mlx_array a, const mlx_stream s) {
   try {

mlx/c/ops.h

@@ -868,6 +868,12 @@ int mlx_scatter_prod(
     const int* axes,
     size_t axes_num,
     const mlx_stream s);
+int mlx_segmented_mm(
+    mlx_array* res,
+    const mlx_array a,
+    const mlx_array b,
+    const mlx_array segments,
+    const mlx_stream s);
 int mlx_sigmoid(mlx_array* res, const mlx_array a, const mlx_stream s);
 int mlx_sign(mlx_array* res, const mlx_array a, const mlx_stream s);
 int mlx_sin(mlx_array* res, const mlx_array a, const mlx_stream s);

mlx/c/random.cpp

@@ -195,6 +195,32 @@ extern "C" int mlx_random_multivariate_normal(
   }
   return 0;
 }
+extern "C" int mlx_random_normal_broadcast(
+    mlx_array* res,
+    const int* shape,
+    size_t shape_num,
+    mlx_dtype dtype,
+    const mlx_array loc /* may be null */,
+    const mlx_array scale /* may be null */,
+    const mlx_array key /* may be null */,
+    const mlx_stream s) {
+  try {
+    mlx_array_set_(
+        *res,
+        mlx::core::random::normal(
+            std::vector<int>(shape, shape + shape_num),
+            mlx_dtype_to_cpp(dtype),
+            (loc.ctx ? std::make_optional(mlx_array_get_(loc)) : std::nullopt),
+            (scale.ctx ? std::make_optional(mlx_array_get_(scale))
+                       : std::nullopt),
+            (key.ctx ? std::make_optional(mlx_array_get_(key)) : std::nullopt),
+            mlx_stream_get_(s)));
+  } catch (std::exception& e) {
+    mlx_error(e.what());
+    return 1;
+  }
+  return 0;
+}
 extern "C" int mlx_random_normal(
     mlx_array* res,
     const int* shape,

mlx/c/random.h

@@ -88,6 +88,15 @@ int mlx_random_multivariate_normal(
     mlx_dtype dtype,
     const mlx_array key /* may be null */,
     const mlx_stream s);
+int mlx_random_normal_broadcast(
+    mlx_array* res,
+    const int* shape,
+    size_t shape_num,
+    mlx_dtype dtype,
+    const mlx_array loc /* may be null */,
+    const mlx_array scale /* may be null */,
+    const mlx_array key /* may be null */,
+    const mlx_stream s);
 int mlx_random_normal(
     mlx_array* res,
     const int* shape,

python/mlxvariants.py

@@ -19,8 +19,6 @@ def _make_variant_suffixes(name, defs, variants):
         print("OVL", file=sys.stderr)
     if name in variants:
         variants = variants[name]
-        for i, d in enumerate(defs):
-            print("OVL", i, _pretty_string_def(d), " -> ", variants[i], file=sys.stderr)
         if len(variants) != len(defs):
             print("function overloads length:", len(defs), file=sys.stderr)
             for i, d in enumerate(defs):
@@ -29,6 +27,8 @@ def _make_variant_suffixes(name, defs, variants):
             for i, v in enumerate(variants):
                 print(i, v, file=sys.stderr)
             raise RuntimeError("function overloads and namings do not match")
+        for i, d in enumerate(defs):
+            print("OVL", i, _pretty_string_def(d), " -> ", variants[i], file=sys.stderr)
         newdefs = []
         for i, d in enumerate(defs):
             v = variants[i]
@@ -112,7 +112,7 @@ def mlx_core_random(name, defs):
         "permutation": ["", "arange"],
         "split": ["num", ""],
         "uniform": ["", None, None, None],
-        "normal": ["", None, None, None],
+        "normal": ["broadcast", "", None, None, None],
     }
     return _make_variant_suffixes(name, defs, variants)