mlx 0.30.7

data/mlx/mlx/fft.cpp

@@ -0,0 +1,262 @@
+// Copyright © 2023 Apple Inc.
+#include <numeric>
+#include <set>
+
+#include "mlx/fft.h"
+#include "mlx/ops.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::fft {
+
+array fft_impl(
+    const array& a,
+    Shape n,
+    const std::vector<int>& axes,
+    bool real,
+    bool inverse,
+    StreamOrDevice s) {
+  if (a.ndim() < 1) {
+    throw std::invalid_argument(
+        "[fftn] Requires array with at least one dimension.");
+  }
+  if (n.size() != axes.size()) {
+    throw std::invalid_argument("[fftn] Shape and axes have different sizes.");
+  }
+  if (axes.empty()) {
+    return a;
+  }
+
+  std::vector<size_t> valid_axes;
+  for (int ax : axes) {
+    valid_axes.push_back(ax < 0 ? ax + a.ndim() : ax);
+  }
+  std::set<int> unique_axes(valid_axes.begin(), valid_axes.end());
+  if (unique_axes.size() != axes.size()) {
+    std::ostringstream msg;
+    msg << "[fftn] Duplicated axis received " << axes;
+    throw std::invalid_argument(msg.str());
+  }
+  if (*unique_axes.begin() < 0 || *unique_axes.rbegin() >= a.ndim()) {
+    std::ostringstream msg;
+    msg << "[fftn] Invalid axis received for array with " << a.ndim()
+        << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // In the following shape manipulations there are three cases to consider:
+  // 1. In a complex to complex transform (fftn / ifftn) the output
+  //    and input shapes are the same.
+  // 2. In a real to complex transform (rfftn) n specifies the input dims
+  //    and the output dims are n[i] / 2 + 1
+  // 3  In a complex to real transform (irfftn) n specifies the output dims
+  //    and the input dims are n[i] / 2 + 1
+
+  if (std::any_of(n.begin(), n.end(), [](auto i) { return i <= 0; })) {
+    std::ostringstream msg;
+    msg << "[fftn] Invalid FFT output size requested " << n;
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto in_shape = a.shape();
+  for (int i = 0; i < valid_axes.size(); ++i) {
+    in_shape[valid_axes[i]] = n[i];
+  }
+  if (real && inverse) {
+    in_shape[valid_axes.back()] = n.back() / 2 + 1;
+  }
+
+  bool any_greater = false;
+  bool any_less = false;
+  for (int i = 0; i < in_shape.size(); ++i) {
+    any_greater |= in_shape[i] > a.shape()[i];
+    any_less |= in_shape[i] < a.shape()[i];
+  }
+
+  auto in = a;
+  if (any_less) {
+    in = slice(in, Shape(in.ndim(), 0), in_shape, s);
+  }
+  if (any_greater) {
+    // Pad with zeros
+    auto tmp = zeros(in_shape, a.dtype(), s);
+    in = slice_update(tmp, in, Shape(in.ndim(), 0), in.shape());
+  }
+
+  auto out_shape = in_shape;
+  if (real) {
+    auto ax = valid_axes.back();
+    out_shape[ax] = inverse ? n.back() : out_shape[ax] / 2 + 1;
+  }
+
+  auto in_type = real && !inverse ? float32 : complex64;
+  auto out_type = real && inverse ? float32 : complex64;
+  return array(
+      out_shape,
+      out_type,
+      std::make_shared<FFT>(to_stream(s), valid_axes, inverse, real),
+      {astype(in, in_type, s)});
+}
+
+array fft_impl(
+    const array& a,
+    const std::vector<int>& axes,
+    bool real,
+    bool inverse,
+    StreamOrDevice s) {
+  Shape n;
+  for (auto ax : axes) {
+    n.push_back(a.shape(ax));
+  }
+  if (real && inverse && a.ndim() > 0) {
+    n.back() = (n.back() - 1) * 2;
+  }
+  return fft_impl(a, n, axes, real, inverse, s);
+}
+
+array fft_impl(const array& a, bool real, bool inverse, StreamOrDevice s) {
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return fft_impl(a, axes, real, inverse, s);
+}
+
+array fftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, n, axes, false, false, s);
+}
+array fftn(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, axes, false, false, s);
+}
+array fftn(const array& a, StreamOrDevice s /* = {} */) {
+  return fft_impl(a, false, false, s);
+}
+
+array ifftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, n, axes, false, true, s);
+}
+array ifftn(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, axes, false, true, s);
+}
+array ifftn(const array& a, StreamOrDevice s /* = {} */) {
+  return fft_impl(a, false, true, s);
+}
+
+array rfftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, n, axes, true, false, s);
+}
+array rfftn(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, axes, true, false, s);
+}
+array rfftn(const array& a, StreamOrDevice s /* = {} */) {
+  return fft_impl(a, true, false, s);
+}
+
+array irfftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, n, axes, true, true, s);
+}
+array irfftn(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  return fft_impl(a, axes, true, true, s);
+}
+
+array irfftn(const array& a, StreamOrDevice s /* = {} */) {
+  return fft_impl(a, true, true, s);
+}
+
+array fftshift(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  if (axes.empty()) {
+    return a;
+  }
+
+  Shape shifts;
+  for (int ax : axes) {
+    // Convert negative axes to positive
+    int axis = ax < 0 ? ax + a.ndim() : ax;
+    if (axis < 0 || axis >= a.ndim()) {
+      std::ostringstream msg;
+      msg << "[fftshift] Invalid axis " << ax << " for array with " << a.ndim()
+          << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    // Match NumPy's implementation
+    shifts.push_back(a.shape(axis) / 2);
+  }
+
+  return roll(a, shifts, axes, s);
+}
+
+array ifftshift(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  if (axes.empty()) {
+    return a;
+  }
+
+  Shape shifts;
+  for (int ax : axes) {
+    // Convert negative axes to positive
+    int axis = ax < 0 ? ax + a.ndim() : ax;
+    if (axis < 0 || axis >= a.ndim()) {
+      std::ostringstream msg;
+      msg << "[ifftshift] Invalid axis " << ax << " for array with " << a.ndim()
+          << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    // Match NumPy's implementation
+    int size = a.shape(axis);
+    shifts.push_back(-(size / 2));
+  }
+
+  return roll(a, shifts, axes, s);
+}
+
+// Default versions that operate on all axes
+array fftshift(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.ndim() < 1) {
+    return a;
+  }
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return fftshift(a, axes, s);
+}
+
+array ifftshift(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.ndim() < 1) {
+    return a;
+  }
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return ifftshift(a, axes, s);
+}
+
+} // namespace mlx::core::fft

data/mlx/mlx/fft.h

@@ -0,0 +1,159 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <variant>
+
+#include "array.h"
+#include "device.h"
+#include "mlx/api.h"
+#include "utils.h"
+
+namespace mlx::core::fft {
+
+/** Compute the n-dimensional Fourier Transform. */
+MLX_API array fftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+MLX_API array
+fftn(const array& a, const std::vector<int>& axes, StreamOrDevice s = {});
+MLX_API array fftn(const array& a, StreamOrDevice s = {});
+
+/** Compute the n-dimensional inverse Fourier Transform. */
+MLX_API array ifftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+MLX_API array
+ifftn(const array& a, const std::vector<int>& axes, StreamOrDevice s = {});
+MLX_API array ifftn(const array& a, StreamOrDevice s = {});
+
+/** Compute the one-dimensional Fourier Transform. */
+inline array fft(const array& a, int n, int axis, StreamOrDevice s = {}) {
+  return fftn(a, {n}, {axis}, s);
+}
+inline array fft(const array& a, int axis = -1, StreamOrDevice s = {}) {
+  return fftn(a, {axis}, s);
+}
+
+/** Compute the one-dimensional inverse Fourier Transform. */
+inline array ifft(const array& a, int n, int axis, StreamOrDevice s = {}) {
+  return ifftn(a, {n}, {axis}, s);
+}
+inline array ifft(const array& a, int axis = -1, StreamOrDevice s = {}) {
+  return ifftn(a, {axis}, s);
+}
+
+/** Compute the two-dimensional Fourier Transform. */
+inline array fft2(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {}) {
+  return fftn(a, n, axes, s);
+}
+inline array fft2(
+    const array& a,
+    const std::vector<int>& axes = {-2, -1},
+    StreamOrDevice s = {}) {
+  return fftn(a, axes, s);
+}
+
+/** Compute the two-dimensional inverse Fourier Transform. */
+inline array ifft2(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {}) {
+  return ifftn(a, n, axes, s);
+}
+inline array ifft2(
+    const array& a,
+    const std::vector<int>& axes = {-2, -1},
+    StreamOrDevice s = {}) {
+  return ifftn(a, axes, s);
+}
+
+/** Compute the n-dimensional Fourier Transform on a real input. */
+MLX_API array rfftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+MLX_API array
+rfftn(const array& a, const std::vector<int>& axes, StreamOrDevice s = {});
+MLX_API array rfftn(const array& a, StreamOrDevice s = {});
+
+/** Compute the n-dimensional inverse of `rfftn`. */
+MLX_API array irfftn(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+MLX_API array
+irfftn(const array& a, const std::vector<int>& axes, StreamOrDevice s = {});
+MLX_API array irfftn(const array& a, StreamOrDevice s = {});
+
+/** Compute the one-dimensional Fourier Transform on a real input. */
+inline array rfft(const array& a, int n, int axis, StreamOrDevice s = {}) {
+  return rfftn(a, {n}, {axis}, s);
+}
+inline array rfft(const array& a, int axis = -1, StreamOrDevice s = {}) {
+  return rfftn(a, {axis}, s);
+}
+/** Compute the one-dimensional inverse of `rfft`. */
+inline array irfft(const array& a, int n, int axis, StreamOrDevice s = {}) {
+  return irfftn(a, {n}, {axis}, s);
+}
+inline array irfft(const array& a, int axis = -1, StreamOrDevice s = {}) {
+  return irfftn(a, {axis}, s);
+}
+
+/** Compute the two-dimensional Fourier Transform on a real input. */
+inline array rfft2(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {}) {
+  return rfftn(a, n, axes, s);
+}
+inline array rfft2(
+    const array& a,
+    const std::vector<int>& axes = {-2, -1},
+    StreamOrDevice s = {}) {
+  return rfftn(a, axes, s);
+}
+
+/** Compute the two-dimensional inverse of `rfft2`. */
+inline array irfft2(
+    const array& a,
+    const Shape& n,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {}) {
+  return irfftn(a, n, axes, s);
+}
+inline array irfft2(
+    const array& a,
+    const std::vector<int>& axes = {-2, -1},
+    StreamOrDevice s = {}) {
+  return irfftn(a, axes, s);
+}
+/** Shift the zero-frequency component to the center of the spectrum. */
+MLX_API array fftshift(const array& a, StreamOrDevice s = {});
+
+/** Shift the zero-frequency component to the center of the spectrum along
+ * specified axes. */
+MLX_API array
+fftshift(const array& a, const std::vector<int>& axes, StreamOrDevice s = {});
+
+/** The inverse of fftshift. */
+MLX_API array ifftshift(const array& a, StreamOrDevice s = {});
+
+/** The inverse of fftshift along specified axes. */
+MLX_API array
+ifftshift(const array& a, const std::vector<int>& axes, StreamOrDevice s = {});
+
+} // namespace mlx::core::fft

data/mlx/mlx/graph_utils.cpp

@@ -0,0 +1,175 @@
+// Copyright © 2023 Apple Inc.
+
+#include <functional>
+#include <optional>
+#include <sstream>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "mlx/graph_utils.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+const std::string& NodeNamer::get_name(const array& x) {
+  auto it = names.find(x.id());
+  if (it == names.end()) {
+    // Get the next name in the sequence
+    // [A, B, ..., Z, AA, AB, ...]
+    std::vector<char> letters;
+    auto var_num = names.size() + 1;
+    while (var_num > 0) {
+      letters.push_back('A' + (var_num - 1) % 26);
+      var_num = (var_num - 1) / 26;
+    }
+    names.emplace(x.id(), std::string(letters.rbegin(), letters.rend()));
+
+    return get_name(x);
+  }
+  return it->second;
+}
+
+void NodeNamer::set_name(const array& x, std::string n) {
+  names[x.id()] = std::move(n);
+}
+
+void depth_first_traversal(
+    std::function<void(array)> callback,
+    const std::vector<array>& outputs) {
+  std::function<void(const array&)> recurse;
+  std::unordered_set<std::uintptr_t> cache;
+  recurse = [&](const array& x) {
+    auto id = x.id();
+    if (cache.find(id) != cache.end()) {
+      return;
+    }
+    cache.insert(id);
+    for (auto& s : x.siblings()) {
+      cache.insert(s.id());
+    }
+    for (auto& in : x.inputs()) {
+      recurse(in);
+    }
+    callback(x);
+  };
+
+  for (auto& o : outputs) {
+    recurse(o);
+  }
+}
+
+void print_graph(
+    std::ostream& os,
+    NodeNamer namer,
+    const std::vector<array>& outputs) {
+  std::vector<array> tape;
+  std::vector<array> inputs;
+
+  depth_first_traversal(
+      [&](const array& x) {
+        if (x.has_primitive()) {
+          tape.push_back(x);
+        } else {
+          inputs.push_back(x);
+        }
+      },
+      outputs);
+
+  auto print_arrs = [&namer, &os](std::vector<array> arrs) {
+    for (auto& arr : arrs) {
+      os << namer.get_name(arr);
+      os << " [" << arr.shape() << ", " << arr.dtype() << "]";
+      if (&arr != &arrs.back()) {
+        os << ", ";
+      }
+    }
+  };
+
+  os << "Inputs: ";
+  print_arrs(inputs);
+  os << "\nOutputs: ";
+  print_arrs(outputs);
+  os << "\n";
+
+  for (auto& arr : tape) {
+    os << arr.primitive().name();
+    os << " ";
+    print_arrs(arr.inputs());
+    os << " -> ";
+    print_arrs(arr.outputs());
+    os << "\n";
+  }
+}
+
+void export_to_dot(
+    std::ostream& os,
+    NodeNamer namer,
+    const std::vector<array>& nodes) {
+  // Perform one DFS to mark arrays as intermediate if they are used as inputs
+  // to other arrays.
+  std::unordered_set<std::uintptr_t> intermediate_set;
+  depth_first_traversal(
+      [&](const array& x) {
+        // No primitive so it is an input
+        if (!x.has_primitive()) {
+          return;
+        }
+
+        for (auto& a : x.inputs()) {
+          intermediate_set.insert(a.id());
+        }
+      },
+      nodes);
+
+  // Now we got everything we need to make the graph. Arrays can be one of 3
+  // things:
+  //  1. Inputs, when they have no primitive ie are evaluated
+  //  2. Intermediates, when they are the intermediate set
+  //  3. Outputs, if they are not inputs and not intermediates
+
+  os << "digraph {" << std::endl;
+
+  depth_first_traversal(
+      [&](const array& x) {
+        if (!x.has_primitive()) {
+          os << "{ rank=source; \"" << namer.get_name(x) << "\"; }"
+             << std::endl;
+          return;
+        }
+
+        // Node for primitive
+        if (x.has_primitive()) {
+          os << "{ ";
+          os << x.primitive_id();
+          os << " [label =\"";
+          os << x.primitive().name();
+          os << "\", shape=rectangle]";
+          os << "; }" << std::endl;
+          // Arrows to primitive's inputs
+          for (auto& a : x.inputs()) {
+            os << '"' << namer.get_name(a) << "\" -> " << x.primitive_id()
+               << std::endl;
+          }
+        }
+
+        // Point outputs to their primitive
+        for (auto& a : x.outputs()) {
+          os << "{ ";
+          if (intermediate_set.find(a.id()) == intermediate_set.end()) {
+            os << "rank=sink; ";
+          }
+          os << '"' << namer.get_name(a);
+          os << "\"; }" << std::endl;
+          if (x.has_primitive()) {
+            os << x.primitive_id() << " -> \"" << namer.get_name(a) << '"'
+               << std::endl;
+          }
+        }
+      },
+      nodes);
+
+  os << "}";
+}
+
+} // namespace mlx::core

data/mlx/mlx/graph_utils.h

@@ -0,0 +1,67 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <unordered_map>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+struct MLX_API NodeNamer {
+  std::unordered_map<std::uintptr_t, std::string> names;
+
+  const std::string& get_name(const array& x);
+  void set_name(const array& x, std::string n);
+};
+
+MLX_API void print_graph(
+    std::ostream& os,
+    NodeNamer namer,
+    const std::vector<array>& outputs);
+
+inline void print_graph(std::ostream& os, const std::vector<array>& outputs) {
+  print_graph(os, NodeNamer{}, outputs);
+}
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+inline void print_graph(std::ostream& os, Arrays&&... outputs) {
+  print_graph(
+      os, NodeNamer{}, std::vector<array>{std::forward<Arrays>(outputs)...});
+}
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+inline void
+print_graph(std::ostream& os, NodeNamer namer, Arrays&&... outputs) {
+  print_graph(
+      os,
+      std::move(namer),
+      std::vector<array>{std::forward<Arrays>(outputs)...});
+}
+
+MLX_API void export_to_dot(
+    std::ostream& os,
+    NodeNamer namer,
+    const std::vector<array>& outputs);
+
+inline void export_to_dot(std::ostream& os, const std::vector<array>& outputs) {
+  export_to_dot(os, NodeNamer{}, outputs);
+}
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+inline void export_to_dot(std::ostream& os, Arrays&&... outputs) {
+  export_to_dot(
+      os, NodeNamer{}, std::vector<array>{std::forward<Arrays>(outputs)...});
+}
+
+template <typename... Arrays, typename = enable_for_arrays_t<Arrays...>>
+inline void
+export_to_dot(std::ostream& os, NodeNamer namer, Arrays&&... outputs) {
+  export_to_dot(
+      os,
+      std::move(namer),
+      std::vector<array>{std::forward<Arrays>(outputs)...});
+}
+
+} // namespace mlx::core

data/mlx/mlx/io/CMakeLists.txt

@@ -0,0 +1,25 @@
+target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/load.cpp)
+
+if(MLX_BUILD_SAFETENSORS)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/safetensors.cpp)
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_safetensors.cpp)
+endif()
+
+if(MLX_BUILD_GGUF)
+  message(STATUS "Downloading gguflib")
+  FetchContent_Declare(
+    gguflib
+    GIT_REPOSITORY https://github.com/antirez/gguf-tools/
+    GIT_TAG 8fa6eb65236618e28fd7710a0fba565f7faa1848)
+  FetchContent_MakeAvailable(gguflib)
+  target_include_directories(mlx
+                             PRIVATE $<BUILD_INTERFACE:${gguflib_SOURCE_DIR}>)
+  add_library(gguflib STATIC ${gguflib_SOURCE_DIR}/fp16.c
+                             ${gguflib_SOURCE_DIR}/gguflib.c)
+  target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:gguflib>)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gguf.cpp
+                             ${CMAKE_CURRENT_SOURCE_DIR}/gguf_quants.cpp)
+else()
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_gguf.cpp)
+endif()