mlx-cpu 0.30.1__py3-none-manylinux_2_35_x86_64.whl

mlx/include/mlx/io/gguf.h

@@ -0,0 +1,20 @@
+// Copyright © 2023-2024 Apple Inc.
+#pragma once
+
+#include "mlx/io.h"
+#include "mlx/primitives.h"
+#include "mlx/transforms.h"
+#include "mlx/utils.h"
+
+extern "C" {
+#include <gguflib.h>
+}
+
+namespace mlx::core {
+
+Shape get_shape(const gguf_tensor& tensor);
+void gguf_load_quantized(
+    std::unordered_map<std::string, array>& a,
+    const gguf_tensor& tensor);
+
+} // namespace mlx::core

mlx/include/mlx/io/load.h

@@ -0,0 +1,175 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <memory>
+#include <sstream>
+
+#include <fcntl.h>
+#ifdef _MSC_VER
+#include <io.h>
+#else
+#include <sys/stat.h>
+#include <unistd.h>
+#endif
+
+#include "mlx/threadpool.h"
+
+// Strictly we need to operate on files in binary mode (to avoid \r getting
+// automatically inserted), but every modern system except for Windows no
+// longer differentiates between binary and text files and for them define
+// the flag as no-op.
+#ifndef O_BINARY
+#define O_BINARY 0
+#endif
+
+namespace mlx::core {
+
+namespace io {
+
+ThreadPool& thread_pool();
+
+class Reader {
+ public:
+  virtual bool is_open() const = 0;
+  virtual bool good() const = 0;
+  virtual size_t tell() = 0; // tellp is non-const in iostream
+  virtual void seek(
+      int64_t off,
+      std::ios_base::seekdir way = std::ios_base::beg) = 0;
+  virtual void read(char* data, size_t n) = 0;
+  virtual void read(char* data, size_t n, size_t offset) = 0;
+  virtual std::string label() const = 0;
+  virtual ~Reader() = default;
+};
+
+class Writer {
+ public:
+  virtual bool is_open() const = 0;
+  virtual bool good() const = 0;
+  virtual size_t tell() = 0;
+  virtual void seek(
+      int64_t off,
+      std::ios_base::seekdir way = std::ios_base::beg) = 0;
+  virtual void write(const char* data, size_t n) = 0;
+  virtual std::string label() const = 0;
+  virtual ~Writer() = default;
+};
+
+class ParallelFileReader : public Reader {
+ public:
+  explicit ParallelFileReader(std::string file_path)
+      : fd_(open(file_path.c_str(), O_RDONLY | O_BINARY)),
+        label_(std::move(file_path)) {}
+
+  ~ParallelFileReader() override {
+    close(fd_);
+  }
+
+  bool is_open() const override {
+    return fd_ > 0;
+  }
+
+  bool good() const override {
+    return is_open();
+  }
+
+  size_t tell() override {
+    return lseek(fd_, 0, SEEK_CUR);
+  }
+
+  // Warning: do not use this function from multiple threads as
+  // it advances the file descriptor
+  void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg)
+      override {
+    if (way == std::ios_base::beg) {
+      lseek(fd_, off, 0);
+    } else {
+      lseek(fd_, off, SEEK_CUR);
+    }
+  }
+
+  // Warning: do not use this function from multiple threads as
+  // it advances the file descriptor
+  void read(char* data, size_t n) override;
+
+  void read(char* data, size_t n, size_t offset) override;
+
+  std::string label() const override {
+    return "file " + label_;
+  }
+
+ private:
+  static constexpr size_t batch_size_ = 1 << 25;
+  static ThreadPool& thread_pool();
+  int fd_;
+  std::string label_;
+};
+
+class FileWriter : public Writer {
+ public:
+  explicit FileWriter() {}
+  explicit FileWriter(std::string file_path)
+      : fd_(open(
+            file_path.c_str(),
+            O_CREAT | O_WRONLY | O_TRUNC | O_BINARY,
+            0644)),
+        label_(std::move(file_path)) {}
+
+  FileWriter(const FileWriter&) = delete;
+  FileWriter& operator=(const FileWriter&) = delete;
+  FileWriter(FileWriter&& other) {
+    std::swap(fd_, other.fd_);
+  }
+
+  ~FileWriter() override {
+    if (fd_ != 0) {
+      close(fd_);
+    }
+  }
+
+  bool is_open() const override {
+    return fd_ >= 0;
+  }
+
+  bool good() const override {
+    return is_open();
+  }
+
+  size_t tell() override {
+    return lseek(fd_, 0, SEEK_CUR);
+  }
+
+  void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg)
+      override {
+    if (way == std::ios_base::beg) {
+      lseek(fd_, off, 0);
+    } else {
+      lseek(fd_, off, SEEK_CUR);
+    }
+  }
+
+  void write(const char* data, size_t n) override {
+    while (n != 0) {
+      auto m = ::write(fd_, data, std::min(n, static_cast<size_t>(INT32_MAX)));
+      if (m <= 0) {
+        std::ostringstream msg;
+        msg << "[write] Unable to write " << n << " bytes to file.";
+        throw std::runtime_error(msg.str());
+      }
+      data += m;
+      n -= m;
+    }
+  }
+
+  std::string label() const override {
+    return "file " + label_;
+  }
+
+ private:
+  int fd_{0};
+  std::string label_;
+};
+
+} // namespace io
+} // namespace mlx::core

mlx/include/mlx/io.h

@@ -0,0 +1,61 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include <unordered_map>
+#include <variant>
+
+#include "mlx/array.h"
+#include "mlx/io/load.h"
+#include "mlx/stream.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+using GGUFMetaData =
+    std::variant<std::monostate, array, std::string, std::vector<std::string>>;
+using GGUFLoad = std::pair<
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, GGUFMetaData>>;
+using SafetensorsLoad = std::pair<
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, std::string>>;
+
+/** Save array to out stream in .npy format */
+void save(std::shared_ptr<io::Writer> out_stream, array a);
+
+/** Save array to file in .npy format */
+void save(std::string file, array a);
+
+/** Load array from reader in .npy format */
+array load(std::shared_ptr<io::Reader> in_stream, StreamOrDevice s = {});
+
+/** Load array from file in .npy format */
+array load(std::string file, StreamOrDevice s = {});
+
+/** Load array map from .safetensors file format */
+SafetensorsLoad load_safetensors(
+    std::shared_ptr<io::Reader> in_stream,
+    StreamOrDevice s = {});
+SafetensorsLoad load_safetensors(
+    const std::string& file,
+    StreamOrDevice s = {});
+
+void save_safetensors(
+    std::shared_ptr<io::Writer> in_stream,
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, std::string> metadata = {});
+void save_safetensors(
+    std::string file,
+    std::unordered_map<std::string, array>,
+    std::unordered_map<std::string, std::string> metadata = {});
+
+/** Load array map and metadata from .gguf file format */
+
+GGUFLoad load_gguf(const std::string& file, StreamOrDevice s = {});
+
+void save_gguf(
+    std::string file,
+    std::unordered_map<std::string, array> array_map,
+    std::unordered_map<std::string, GGUFMetaData> meta_data = {});
+
+} // namespace mlx::core

mlx/include/mlx/linalg.h

@@ -0,0 +1,111 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <optional>
+
+#include "mlx/array.h"
+#include "mlx/device.h"
+#include "mlx/ops.h"
+#include "mlx/stream.h"
+
+namespace mlx::core::linalg {
+
+/**
+ * Compute vector or matrix norms.
+ *
+ * - If axis and ord are both unspecified, computes the 2-norm of flatten(x).
+ * - If axis is not provided but ord is, then x must be either 1D or 2D.
+ * - If axis is provided, but ord is not, then the 2-norm (or Frobenius norm
+ *   for matrices) is computed along the given axes. At most 2 axes can be
+ *   specified.
+ * - If both axis and ord are provided, then the corresponding matrix or vector
+ *   norm is computed. At most 2 axes can be specified.
+ */
+array norm(
+    const array& a,
+    const double ord,
+    const std::optional<std::vector<int>>& axis = std::nullopt,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+inline array norm(
+    const array& a,
+    const double ord,
+    int axis,
+    bool keepdims = false,
+    StreamOrDevice s = {}) {
+  return norm(a, ord, std::vector<int>{axis}, keepdims, s);
+}
+array norm(
+    const array& a,
+    const std::string& ord,
+    const std::optional<std::vector<int>>& axis = std::nullopt,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+inline array norm(
+    const array& a,
+    const std::string& ord,
+    int axis,
+    bool keepdims = false,
+    StreamOrDevice s = {}) {
+  return norm(a, ord, std::vector<int>{axis}, keepdims, s);
+}
+array norm(
+    const array& a,
+    const std::optional<std::vector<int>>& axis = std::nullopt,
+    bool keepdims = false,
+    StreamOrDevice s = {});
+inline array
+norm(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}) {
+  return norm(a, std::vector<int>{axis}, keepdims, s);
+}
+
+std::pair<array, array> qr(const array& a, StreamOrDevice s = {});
+
+std::vector<array>
+svd(const array& a, bool compute_uv, StreamOrDevice s /* = {} */);
+inline std::vector<array> svd(const array& a, StreamOrDevice s = {}) {
+  return svd(a, true, s);
+}
+
+array inv(const array& a, StreamOrDevice s = {});
+
+array tri_inv(const array& a, bool upper = false, StreamOrDevice s = {});
+
+array cholesky(const array& a, bool upper = false, StreamOrDevice s = {});
+
+array pinv(const array& a, StreamOrDevice s = {});
+
+array cholesky_inv(const array& a, bool upper = false, StreamOrDevice s = {});
+
+std::vector<array> lu(const array& a, StreamOrDevice s = {});
+
+std::pair<array, array> lu_factor(const array& a, StreamOrDevice s = {});
+
+array solve(const array& a, const array& b, StreamOrDevice s = {});
+
+array solve_triangular(
+    const array& a,
+    const array& b,
+    bool upper = false,
+    StreamOrDevice s = {});
+
+/**
+ * Compute the cross product of two arrays along the given axis.
+ */
+array cross(
+    const array& a,
+    const array& b,
+    int axis = -1,
+    StreamOrDevice s = {});
+
+std::pair<array, array> eig(const array& a, StreamOrDevice s = {});
+
+array eigvals(const array& a, StreamOrDevice s = {});
+
+array eigvalsh(const array& a, std::string UPLO = "L", StreamOrDevice s = {});
+
+std::pair<array, array>
+eigh(const array& a, std::string UPLO = "L", StreamOrDevice s = {});
+
+} // namespace mlx::core::linalg

mlx/include/mlx/memory.h

@@ -0,0 +1,78 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <cstdlib>
+
+namespace mlx::core {
+
+/* Get the actively used memory in bytes.
+ *
+ * Note, this will not always match memory use reported by the system because
+ * it does not include cached memory buffers.
+ * */
+size_t get_active_memory();
+
+/* Get the peak amount of used memory in bytes.
+ *
+ * The maximum memory used recorded from the beginning of the program
+ * execution or since the last call to reset_peak_memory.
+ * */
+size_t get_peak_memory();
+
+/* Reset the peak memory to zero.
+ * */
+void reset_peak_memory();
+
+/* Get the cache size in bytes.
+ *
+ * The cache includes memory not currently used that has not been returned
+ * to the system allocator.
+ * */
+size_t get_cache_memory();
+
+/* Set the memory limit.
+ * The memory limit is a guideline for the maximum amount of memory to use
+ * during graph evaluation. If the memory limit is exceeded and there is no
+ * more RAM (including swap when available) allocations will result in an
+ * exception.
+ *
+ * When Metal is available the memory limit defaults to 1.5 times the maximum
+ * recommended working set size reported by the device.
+ *
+ * Returns the previous memory limit.
+ * */
+size_t set_memory_limit(size_t limit);
+
+/* Get the current memory limit. */
+size_t get_memory_limit();
+
+/* Set the cache limit.
+ * If using more than the given limit, free memory will be reclaimed
+ * from the cache on the next allocation. To disable the cache,
+ * set the limit to 0.
+ *
+ * The cache limit defaults to the memory limit.
+ *
+ * Returns the previous cache limit.
+ * */
+size_t set_cache_limit(size_t limit);
+
+/* Clear the memory cache. */
+void clear_cache();
+
+/* Set the wired size limit.
+ *
+ * Note, this function is only useful when using the Metal backend with
+ * macOS 15.0 or higher.
+ *
+ * The wired limit is the total size in bytes of memory that will be kept
+ * resident. The default value is ``0``.
+ *
+ * Setting a wired limit larger than system wired limit is an error.
+ *
+ * Returns the previous wired limit.
+ * */
+size_t set_wired_limit(size_t limit);
+
+} // namespace mlx::core

mlx/include/mlx/mlx.h

@@ -0,0 +1,25 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/backend/cuda/cuda.h"
+#include "mlx/backend/gpu/available.h"
+#include "mlx/backend/metal/metal.h"
+#include "mlx/compile.h"
+#include "mlx/device.h"
+#include "mlx/distributed/distributed.h"
+#include "mlx/distributed/ops.h"
+#include "mlx/einsum.h"
+#include "mlx/export.h"
+#include "mlx/fast.h"
+#include "mlx/fft.h"
+#include "mlx/io.h"
+#include "mlx/linalg.h"
+#include "mlx/memory.h"
+#include "mlx/ops.h"
+#include "mlx/random.h"
+#include "mlx/stream.h"
+#include "mlx/transforms.h"
+#include "mlx/utils.h"
+#include "mlx/version.h"