mlx 1.0.0

data/mlx/python/src/trees.cpp

@@ -0,0 +1,311 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "python/src/trees.h"
+
+template <typename T, typename U, typename V>
+void validate_subtrees(const std::vector<nb::object>& subtrees) {
+  int len = nb::cast<T>(subtrees[0]).size();
+  for (auto& subtree : subtrees) {
+    if ((nb::isinstance<T>(subtree) && nb::cast<T>(subtree).size() != len) ||
+        nb::isinstance<U>(subtree) || nb::isinstance<V>(subtree)) {
+      throw std::invalid_argument(
+          "[tree_map] Additional input tree is not a valid prefix of the first tree.");
+    }
+  }
+}
+
+nb::object tree_map(
+    const std::vector<nb::object>& trees,
+    std::function<nb::object(const std::vector<nb::object>&)> transform) {
+  std::function<nb::object(const std::vector<nb::object>&)> recurse;
+
+  recurse = [&](const std::vector<nb::object>& subtrees) {
+    if (nb::isinstance<nb::list>(subtrees[0])) {
+      nb::list l;
+      std::vector<nb::object> items(subtrees.size());
+      validate_subtrees<nb::list, nb::tuple, nb::dict>(subtrees);
+      for (int i = 0; i < nb::cast<nb::list>(subtrees[0]).size(); ++i) {
+        for (int j = 0; j < subtrees.size(); ++j) {
+          if (nb::isinstance<nb::list>(subtrees[j])) {
+            items[j] = nb::cast<nb::list>(subtrees[j])[i];
+          } else {
+            items[j] = subtrees[j];
+          }
+        }
+        l.append(recurse(items));
+      }
+      return nb::cast<nb::object>(l);
+    } else if (nb::isinstance<nb::tuple>(subtrees[0])) {
+      //  Check the rest of the subtrees
+      std::vector<nb::object> items(subtrees.size());
+      int len = nb::cast<nb::tuple>(subtrees[0]).size();
+      nb::list l;
+      validate_subtrees<nb::tuple, nb::list, nb::dict>(subtrees);
+      auto type = subtrees[0].type();
+      for (int i = 0; i < len; ++i) {
+        for (int j = 0; j < subtrees.size(); ++j) {
+          if (nb::isinstance<nb::tuple>(subtrees[j])) {
+            items[j] = nb::cast<nb::tuple>(subtrees[j])[i];
+          } else {
+            items[j] = subtrees[j];
+          }
+        }
+        l.append(recurse(items));
+      }
+      if (PyTuple_CheckExact(subtrees[0].ptr())) {
+        return nb::cast<nb::object>(nb::tuple(l));
+      }
+      return nb::hasattr(type, "_fields") ? type(*l) : type(l);
+    } else if (nb::isinstance<nb::dict>(subtrees[0])) {
+      std::vector<nb::object> items(subtrees.size());
+      validate_subtrees<nb::dict, nb::list, nb::tuple>(subtrees);
+      nb::dict d;
+      for (auto item : nb::cast<nb::dict>(subtrees[0])) {
+        for (int j = 0; j < subtrees.size(); ++j) {
+          if (nb::isinstance<nb::dict>(subtrees[j])) {
+            auto subdict = nb::cast<nb::dict>(subtrees[j]);
+            if (!subdict.contains(item.first)) {
+              throw std::invalid_argument(
+                  "[tree_map] Tree is not a valid prefix tree of the first tree.");
+            }
+            items[j] = subdict[item.first];
+          } else {
+            items[j] = subtrees[j];
+          }
+        }
+        d[item.first] = recurse(items);
+      }
+      return nb::cast<nb::object>(d);
+    } else {
+      return transform(subtrees);
+    }
+  };
+  return recurse(trees);
+}
+
+nb::object tree_map(
+    nb::object tree,
+    std::function<nb::object(nb::handle)> transform) {
+  return tree_map({tree}, [&](std::vector<nb::object> inputs) {
+    return transform(inputs[0]);
+  });
+}
+
+void tree_visit(
+    const std::vector<nb::object>& trees,
+    std::function<void(const std::vector<nb::object>&)> visitor) {
+  std::function<void(const std::vector<nb::object>&)> recurse;
+
+  recurse = [&](const std::vector<nb::object>& subtrees) {
+    if (nb::isinstance<nb::list>(subtrees[0])) {
+      std::vector<nb::object> items(subtrees.size());
+      validate_subtrees<nb::list, nb::tuple, nb::dict>(subtrees);
+      for (int i = 0; i < nb::cast<nb::list>(subtrees[0]).size(); ++i) {
+        for (int j = 0; j < subtrees.size(); ++j) {
+          if (nb::isinstance<nb::list>(subtrees[j])) {
+            items[j] = nb::cast<nb::list>(subtrees[j])[i];
+          } else {
+            items[j] = subtrees[j];
+          }
+        }
+        recurse(items);
+      }
+    } else if (nb::isinstance<nb::tuple>(subtrees[0])) {
+      //  Check the rest of the subtrees
+      std::vector<nb::object> items(subtrees.size());
+      int len = nb::cast<nb::tuple>(subtrees[0]).size();
+      validate_subtrees<nb::tuple, nb::list, nb::dict>(subtrees);
+      for (int i = 0; i < len; ++i) {
+        for (int j = 0; j < subtrees.size(); ++j) {
+          if (nb::isinstance<nb::tuple>(subtrees[j])) {
+            items[j] = nb::cast<nb::tuple>(subtrees[j])[i];
+          } else {
+            items[j] = subtrees[j];
+          }
+        }
+        recurse(items);
+      }
+    } else if (nb::isinstance<nb::dict>(subtrees[0])) {
+      std::vector<nb::object> items(subtrees.size());
+      validate_subtrees<nb::dict, nb::list, nb::tuple>(subtrees);
+      for (auto item : nb::cast<nb::dict>(subtrees[0])) {
+        for (int j = 0; j < subtrees.size(); ++j) {
+          if (nb::isinstance<nb::dict>(subtrees[j])) {
+            auto subdict = nb::cast<nb::dict>(subtrees[j]);
+            if (!subdict.contains(item.first)) {
+              throw std::invalid_argument(
+                  "[tree_visit] Tree is not a valid prefix tree of the first tree.");
+            }
+            items[j] = subdict[item.first];
+          } else {
+            items[j] = subtrees[j];
+          }
+        }
+        recurse(items);
+      }
+    } else {
+      visitor(subtrees);
+    }
+  };
+  return recurse(trees);
+}
+
+void tree_visit(nb::handle tree, std::function<void(nb::handle)> visitor) {
+  std::function<void(nb::handle)> recurse;
+  recurse = [&](nb::handle subtree) {
+    if (nb::isinstance<nb::list>(subtree) ||
+        nb::isinstance<nb::tuple>(subtree)) {
+      for (auto item : subtree) {
+        recurse(item);
+      }
+    } else if (nb::isinstance<nb::dict>(subtree)) {
+      for (auto item : nb::cast<nb::dict>(subtree)) {
+        recurse(item.second);
+      }
+    } else {
+      visitor(subtree);
+    }
+  };
+
+  recurse(tree);
+}
+
+void tree_visit_update(
+    nb::object tree,
+    std::function<nb::object(nb::handle)> visitor) {
+  std::function<nb::object(nb::handle)> recurse;
+  recurse = [&](nb::handle subtree) {
+    if (nb::isinstance<nb::list>(subtree)) {
+      auto l = nb::cast<nb::list>(subtree);
+      for (int i = 0; i < l.size(); ++i) {
+        l[i] = recurse(l[i]);
+      }
+      return nb::cast<nb::object>(l);
+    } else if (nb::isinstance<nb::tuple>(subtree)) {
+      auto type = subtree.type();
+      nb::list l(subtree);
+      for (int i = 0; i < l.size(); ++i) {
+        l[i] = recurse(l[i]);
+      }
+      if (PyTuple_CheckExact(subtree.ptr())) {
+        return nb::cast<nb::object>(nb::tuple(l));
+      }
+      return nb::hasattr(type, "_fields") ? type(*l) : type(l);
+    } else if (nb::isinstance<nb::dict>(subtree)) {
+      auto d = nb::cast<nb::dict>(subtree);
+      for (auto item : d) {
+        d[item.first] = recurse(item.second);
+      }
+      return nb::cast<nb::object>(d);
+    } else if (nb::isinstance<mx::array>(subtree)) {
+      return visitor(subtree);
+    } else {
+      return nb::cast<nb::object>(subtree);
+    }
+  };
+  recurse(tree);
+}
+
+// Fill a pytree (recursive dict or list of dict or list)
+// in place with the given arrays
+// Non dict or list nodes are ignored
+void tree_fill(nb::object& tree, const std::vector<mx::array>& values) {
+  size_t index = 0;
+  tree_visit_update(
+      tree, [&](nb::handle node) { return nb::cast(values[index++]); });
+}
+
+// Replace all the arrays from the src values with the dst values in the tree
+void tree_replace(
+    nb::object& tree,
+    const std::vector<mx::array>& src,
+    const std::vector<mx::array>& dst) {
+  std::unordered_map<uintptr_t, mx::array> src_to_dst;
+  for (int i = 0; i < src.size(); ++i) {
+    src_to_dst.insert({src[i].id(), dst[i]});
+  }
+  tree_visit_update(tree, [&](nb::handle node) {
+    auto arr = nb::cast<mx::array>(node);
+    if (auto it = src_to_dst.find(arr.id()); it != src_to_dst.end()) {
+      return nb::cast(it->second);
+    }
+    return nb::cast(arr);
+  });
+}
+
+std::vector<mx::array> tree_flatten(nb::handle tree, bool strict /* = true */) {
+  std::vector<mx::array> flat_tree;
+
+  tree_visit(tree, [&](nb::handle obj) {
+    if (nb::isinstance<mx::array>(obj)) {
+      flat_tree.push_back(nb::cast<mx::array>(obj));
+    } else if (strict) {
+      throw std::invalid_argument(
+          "[tree_flatten] The argument should contain only arrays");
+    }
+  });
+
+  return flat_tree;
+}
+
+nb::object tree_unflatten(
+    nb::object tree,
+    const std::vector<mx::array>& values,
+    int index /* = 0 */) {
+  return tree_map(tree, [&](nb::handle obj) {
+    if (nb::isinstance<mx::array>(obj)) {
+      return nb::cast(values[index++]);
+    } else {
+      return nb::cast<nb::object>(obj);
+    }
+  });
+}
+
+nb::object structure_sentinel() {
+  static nb::object sentinel;
+
+  if (sentinel.ptr() == nullptr) {
+    sentinel = nb::capsule(&sentinel);
+    // probably not needed but this should make certain that we won't ever
+    // delete the sentinel
+    sentinel.inc_ref();
+  }
+
+  return sentinel;
+}
+
+std::pair<std::vector<mx::array>, nb::object> tree_flatten_with_structure(
+    nb::object tree,
+    bool strict /* = true */) {
+  auto sentinel = structure_sentinel();
+  std::vector<mx::array> flat_tree;
+  auto structure = tree_map(
+      tree,
+      [&flat_tree, sentinel = std::move(sentinel), strict](nb::handle obj) {
+        if (nb::isinstance<mx::array>(obj)) {
+          flat_tree.push_back(nb::cast<mx::array>(obj));
+          return sentinel;
+        } else if (!strict) {
+          return nb::cast<nb::object>(obj);
+        } else {
+          throw std::invalid_argument(
+              "[tree_flatten] The argument should contain only arrays");
+        }
+      });
+
+  return {flat_tree, structure};
+}
+
+nb::object tree_unflatten_from_structure(
+    nb::object structure,
+    const std::vector<mx::array>& values,
+    int index /* = 0 */) {
+  auto sentinel = structure_sentinel();
+  return tree_map(structure, [&](nb::handle obj) {
+    if (obj.is(sentinel)) {
+      return nb::cast(values[index++]);
+    } else {
+      return nb::cast<nb::object>(obj);
+    }
+  });
+}

data/mlx/python/src/trees.h

@@ -0,0 +1,62 @@
+// Copyright © 2023-2024 Apple Inc.
+#pragma once
+#include <nanobind/nanobind.h>
+
+#include "mlx/array.h"
+
+namespace mx = mlx::core;
+namespace nb = nanobind;
+
+void tree_visit(
+    const std::vector<nb::object>& trees,
+    std::function<void(const std::vector<nb::object>&)> visitor);
+void tree_visit(nb::handle tree, std::function<void(nb::handle)> visitor);
+
+nb::object tree_map(
+    const std::vector<nb::object>& trees,
+    std::function<nb::object(const std::vector<nb::object>&)> transform);
+
+nb::object tree_map(
+    nb::object tree,
+    std::function<nb::object(nb::handle)> transform);
+
+void tree_visit_update(
+    nb::object tree,
+    std::function<nb::object(nb::handle)> visitor);
+
+/**
+ * Fill a pytree (recursive dict or list of dict or list) in place with the
+ * given arrays. */
+void tree_fill(nb::object& tree, const std::vector<mx::array>& values);
+
+/**
+ * Replace all the arrays from the src values with the dst values in the
+ * tree.
+ */
+void tree_replace(
+    nb::object& tree,
+    const std::vector<mx::array>& src,
+    const std::vector<mx::array>& dst);
+
+/**
+ * Flatten a tree into a vector of arrays. If strict is true, then the
+ * function will throw if the tree contains a leaf which is not an array.
+ */
+std::vector<mx::array> tree_flatten(nb::handle tree, bool strict = true);
+
+/**
+ * Unflatten a tree from a vector of arrays.
+ */
+nb::object tree_unflatten(
+    nb::object tree,
+    const std::vector<mx::array>& values,
+    int index = 0);
+
+std::pair<std::vector<mx::array>, nb::object> tree_flatten_with_structure(
+    nb::object tree,
+    bool strict = true);
+
+nb::object tree_unflatten_from_structure(
+    nb::object structure,
+    const std::vector<mx::array>& values,
+    int index = 0);

data/mlx/python/src/utils.cpp

@@ -0,0 +1,98 @@
+// Copyright © 2024 Apple Inc.
+
+#include "python/src/utils.h"
+#include "mlx/ops.h"
+#include "mlx/utils.h"
+#include "python/src/convert.h"
+
+mx::array to_array(
+    const ScalarOrArray& v,
+    std::optional<mx::Dtype> dtype /* = std::nullopt */) {
+  if (auto pv = std::get_if<nb::bool_>(&v); pv) {
+    return mx::array(nb::cast<bool>(*pv), dtype.value_or(mx::bool_));
+  } else if (auto pv = std::get_if<nb::int_>(&v); pv) {
+    auto val = nb::cast<int64_t>(*pv);
+    auto default_type = (val > std::numeric_limits<int>::max() ||
+                         val < std::numeric_limits<int>::min())
+        ? mx::int64
+        : mx::int32;
+    auto out_t = dtype.value_or(default_type);
+    if (mx::issubdtype(out_t, mx::integer) && out_t.size() < 8) {
+      auto info = mx::iinfo(out_t);
+      if (val < info.min || val > static_cast<int64_t>(info.max)) {
+        std::ostringstream msg;
+        msg << "Converting " << val << " to " << out_t
+            << " would result in overflow.";
+        throw std::invalid_argument(msg.str());
+      }
+    }
+
+    // bool_ is an exception and is always promoted
+    return mx::array(val, (out_t == mx::bool_) ? mx::int32 : out_t);
+  } else if (auto pv = std::get_if<nb::float_>(&v); pv) {
+    auto out_t = dtype.value_or(mx::float32);
+    return mx::array(
+        nb::cast<float>(*pv),
+        mx::issubdtype(out_t, mx::floating) ? out_t : mx::float32);
+  } else if (auto pv = std::get_if<std::complex<float>>(&v); pv) {
+    return mx::array(static_cast<mx::complex64_t>(*pv), mx::complex64);
+  } else if (auto pv = std::get_if<mx::array>(&v); pv) {
+    return *pv;
+  } else if (auto pv = std::get_if<
+                 nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>>(&v);
+             pv) {
+    return nd_array_to_mlx(*pv, dtype);
+  } else {
+    return to_array_with_accessor(std::get<ArrayLike>(v).obj);
+  }
+}
+
+std::pair<mx::array, mx::array> to_arrays(
+    const ScalarOrArray& a,
+    const ScalarOrArray& b) {
+  // Four cases:
+  // - If both a and b are arrays leave their types alone
+  // - If a is an array but b is not, treat b as a weak python type
+  // - If b is an array but a is not, treat a as a weak python type
+  // - If neither is an array convert to arrays but leave their types alone
+  auto is_mlx_array = [](const ScalarOrArray& x) {
+    return std::holds_alternative<mx::array>(x) ||
+        std::holds_alternative<ArrayLike>(x) &&
+        nb::hasattr(std::get<ArrayLike>(x).obj, "__mlx_array__");
+  };
+  auto get_mlx_array = [](const ScalarOrArray& x) {
+    if (auto px = std::get_if<mx::array>(&x); px) {
+      return *px;
+    } else {
+      return nb::cast<mx::array>(
+          std::get<ArrayLike>(x).obj.attr("__mlx_array__"));
+    }
+  };
+
+  if (is_mlx_array(a)) {
+    auto arr_a = get_mlx_array(a);
+    if (is_mlx_array(b)) {
+      auto arr_b = get_mlx_array(b);
+      return {arr_a, arr_b};
+    }
+    return {arr_a, to_array(b, arr_a.dtype())};
+  } else if (is_mlx_array(b)) {
+    auto arr_b = get_mlx_array(b);
+    return {to_array(a, arr_b.dtype()), arr_b};
+  } else {
+    return {to_array(a), to_array(b)};
+  }
+}
+
+mx::array to_array_with_accessor(nb::object obj) {
+  if (nb::isinstance<mx::array>(obj)) {
+    return nb::cast<mx::array>(obj);
+  } else if (nb::hasattr(obj, "__mlx_array__")) {
+    return nb::cast<mx::array>(obj.attr("__mlx_array__")());
+  } else {
+    std::ostringstream msg;
+    msg << "Invalid type " << nb::type_name(obj.type()).c_str()
+        << " received in array initialization.";
+    throw std::invalid_argument(msg.str());
+  }
+}

data/mlx/python/src/utils.h

@@ -0,0 +1,78 @@
+// Copyright © 2023-2024 Apple Inc.
+#pragma once
+#include <numeric>
+#include <optional>
+#include <string>
+#include <variant>
+
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+#include <nanobind/stl/complex.h>
+#include <nanobind/stl/variant.h>
+
+#include "mlx/array.h"
+#include "python/src/convert.h"
+
+namespace mx = mlx::core;
+namespace nb = nanobind;
+
+using IntOrVec = std::variant<std::monostate, int, std::vector<int>>;
+using ScalarOrArray = std::variant<
+    nb::bool_,
+    nb::int_,
+    nb::float_,
+    // Must be above ndarray
+    mx::array,
+    // Must be above complex
+    nb::ndarray<nb::ro, nb::c_contig, nb::device::cpu>,
+    std::complex<float>,
+    ArrayLike>;
+
+inline std::vector<int> get_reduce_axes(const IntOrVec& v, int dims) {
+  std::vector<int> axes;
+  if (std::holds_alternative<std::monostate>(v)) {
+    axes.resize(dims);
+    std::iota(axes.begin(), axes.end(), 0);
+  } else if (auto pv = std::get_if<int>(&v); pv) {
+    axes.push_back(*pv);
+  } else {
+    axes = std::get<std::vector<int>>(v);
+  }
+  return axes;
+}
+
+inline bool is_comparable_with_array(const ScalarOrArray& v) {
+  // Checks if the value can be compared to an array (or is already an
+  // mlx array)
+  if (auto pv = std::get_if<ArrayLike>(&v); pv) {
+    auto obj = (*pv).obj;
+    return nb::isinstance<mx::array>(obj) || nb::hasattr(obj, "__mlx_array__");
+  } else {
+    // If it's not an object, it's a scalar (nb::int_, nb::float_, etc.)
+    // and can be compared to an array
+    return true;
+  }
+}
+
+inline nb::handle get_handle_of_object(const ScalarOrArray& v) {
+  return std::get<ArrayLike>(v).obj.ptr();
+}
+
+inline void throw_invalid_operation(
+    const std::string& operation,
+    const ScalarOrArray operand) {
+  std::ostringstream msg;
+  msg << "Cannot perform " << operation << " on an mlx.core.array and "
+      << nb::type_name(get_handle_of_object(operand).type()).c_str();
+  throw std::invalid_argument(msg.str());
+}
+
+mx::array to_array(
+    const ScalarOrArray& v,
+    std::optional<mx::Dtype> dtype = std::nullopt);
+
+std::pair<mx::array, mx::array> to_arrays(
+    const ScalarOrArray& a,
+    const ScalarOrArray& b);
+
+mx::array to_array_with_accessor(nb::object obj);

data/mlx/python/tests/__main__.py

@@ -0,0 +1,5 @@
+from . import mlx_tests
+
+__unittest = True
+
+mlx_tests.MLXTestRunner(module=None)

data/mlx/python/tests/cuda_skip.py

@@ -0,0 +1,62 @@
+cuda_skip = {
+    "TestLayers.test_quantized_embedding",
+    # Block masked matmul NYI
+    "TestBlas.test_block_masked_matmul",
+    # Gather matmul NYI
+    "TestBlas.test_gather_matmul",
+    "TestBlas.test_gather_matmul_grad",
+    "TestBlas.test_gather_mm_sorted_vjp",
+    # Segmented matmul NYI
+    "TestBlas.test_segmented_mm",
+    # Hadamard NYI
+    "TestOps.test_hadamard",
+    "TestOps.test_hadamard_grad_vmap",
+    # FFTs NYI
+    "TestFFT.test_fft",
+    "TestFFT.test_fft_big_powers_of_two",
+    "TestFFT.test_fft_contiguity",
+    "TestFFT.test_fft_exhaustive",
+    "TestFFT.test_fft_grads",
+    "TestFFT.test_fft_into_ifft",
+    "TestFFT.test_fft_large_numbers",
+    "TestFFT.test_fft_shared_mem",
+    "TestFFT.test_fftn",
+    # Lapack ops NYI
+    "TestLinalg.test_cholesky",
+    "TestLinalg.test_cholesky_inv",
+    "TestLinalg.test_eig",
+    "TestLinalg.test_eigh",
+    "TestLinalg.test_inverse",
+    "TestVmap.test_vmap_inverse",
+    "TestLinalg.test_lu",
+    "TestLinalg.test_lu_factor",
+    "TestLinalg.test_pseudo_inverse",
+    "TestLinalg.test_qr_factorization",
+    "TestInit.test_orthogonal",
+    "TestLinalg.test_svd_decomposition",
+    "TestVmap.test_vmap_svd",
+    "TestLinalg.test_tri_inverse",
+    # Quantization NYI
+    "TestQuantized.test_gather_matmul_grad",
+    "TestQuantized.test_gather_qmm",
+    "TestQuantized.test_gather_qmm_sorted",
+    "TestQuantized.test_gather_qmm_grad",
+    "TestQuantized.test_non_multiples",
+    "TestQuantized.test_qmm",
+    "TestQuantized.test_qmm_jvp",
+    "TestQuantized.test_qmm_shapes",
+    "TestQuantized.test_qmm_vjp",
+    "TestQuantized.test_qmv",
+    "TestQuantized.test_fp_qvm",
+    "TestQuantized.test_qvm",
+    "TestQuantized.test_qvm_splitk",
+    "TestQuantized.test_qmv_small_non_multiples",
+    "TestQuantized.test_small_matrix",
+    "TestQuantized.test_throw",
+    "TestQuantized.test_vjp_scales_biases",
+    "TestExportImport.test_export_quantized_model",
+    # Masked scatter
+    "TestOps.test_masked_scatter",
+    "TestVmap.test_vmap_masked_scatter",
+    "TestArray.test_setitem_with_boolean_mask",
+}