mlx 0.30.7.2 → 0.30.7.6

data/{mlx → submodules/mlx}/mlx/stream.h

@@ -2,6 +2,8 @@
 
 #pragma once
 
+#include <vector>
+
 #include "mlx/api.h"
 #include "mlx/device.h"
 
@@ -25,6 +27,9 @@ MLX_API Stream new_stream(Device d);
 /** Get the stream with the given index. */
 MLX_API Stream get_stream(int index);
 
+/** Get all available streams. */
+MLX_API std::vector<Stream> get_streams();
+
 inline bool operator==(const Stream& lhs, const Stream& rhs) {
   return lhs.index == rhs.index;
 }

data/submodules/mlx-onnx/CMakeLists.txt

@@ -0,0 +1,159 @@
+cmake_minimum_required(VERSION 3.25)
+
+project(mlx_onnx VERSION 0.30.7.1 LANGUAGES C CXX)
+
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+
+option(MLX_ONNX_USE_EXTERNAL_MLX "Build against an externally provided MLX install" OFF)
+option(MLX_ONNX_BUILD_PYTHON_BINDINGS "Build Python IR bindings" OFF)
+option(MLX_ONNX_INSTALL_CPP_ARTIFACTS "Install C++ library and headers" ON)
+
+include(FetchContent)
+
+if(MLX_ONNX_USE_EXTERNAL_MLX AND MLX_ONNX_BUILD_PYTHON_BINDINGS)
+  message(
+    FATAL_ERROR
+      "MLX_ONNX_BUILD_PYTHON_BINDINGS requires bundled mlx sources; set MLX_ONNX_USE_EXTERNAL_MLX=OFF")
+endif()
+
+if(MLX_ONNX_USE_EXTERNAL_MLX)
+  set(MLX_ONNX_EXTERNAL_MLX_INCLUDE_DIR "" CACHE PATH "Path to MLX include root")
+  set(MLX_ONNX_EXTERNAL_MLX_LIB_DIR "" CACHE PATH "Path to MLX library directory")
+
+  if(MLX_ONNX_EXTERNAL_MLX_INCLUDE_DIR STREQUAL "")
+    message(FATAL_ERROR "MLX_ONNX_EXTERNAL_MLX_INCLUDE_DIR must be set when MLX_ONNX_USE_EXTERNAL_MLX=ON")
+  endif()
+  if(MLX_ONNX_EXTERNAL_MLX_LIB_DIR STREQUAL "")
+    message(FATAL_ERROR "MLX_ONNX_EXTERNAL_MLX_LIB_DIR must be set when MLX_ONNX_USE_EXTERNAL_MLX=ON")
+  endif()
+
+  find_library(
+    MLX_EXTERNAL_LIBRARY
+    NAMES mlx
+    PATHS ${MLX_ONNX_EXTERNAL_MLX_LIB_DIR}
+    NO_DEFAULT_PATH)
+
+  if(NOT MLX_EXTERNAL_LIBRARY)
+    message(FATAL_ERROR "Could not find libmlx in ${MLX_ONNX_EXTERNAL_MLX_LIB_DIR}")
+  endif()
+
+  add_library(mlx SHARED IMPORTED GLOBAL)
+  set_target_properties(
+    mlx
+    PROPERTIES
+      IMPORTED_LOCATION ${MLX_EXTERNAL_LIBRARY}
+      INTERFACE_INCLUDE_DIRECTORIES ${MLX_ONNX_EXTERNAL_MLX_INCLUDE_DIR})
+else()
+  set(MLX_BUILD_TESTS OFF CACHE BOOL "" FORCE)
+  set(MLX_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE)
+  set(MLX_BUILD_BENCHMARKS OFF CACHE BOOL "" FORCE)
+  if(MLX_ONNX_BUILD_PYTHON_BINDINGS)
+    set(MLX_BUILD_PYTHON_BINDINGS ON CACHE BOOL "" FORCE)
+  else()
+    set(MLX_BUILD_PYTHON_BINDINGS OFF CACHE BOOL "" FORCE)
+  endif()
+  set(MLX_BUILD_PYTHON_STUBS OFF CACHE BOOL "" FORCE)
+  set(MLX_BUILD_GGUF OFF CACHE BOOL "" FORCE)
+  set(MLX_BUILD_SAFETENSORS OFF CACHE BOOL "" FORCE)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/mlx)
+endif()
+
+if(NOT TARGET nlohmann_json::nlohmann_json)
+  FetchContent_Declare(
+    nlohmann_json
+    GIT_REPOSITORY https://github.com/nlohmann/json.git
+    GIT_TAG v3.11.3
+    EXCLUDE_FROM_ALL)
+  FetchContent_MakeAvailable(nlohmann_json)
+endif()
+
+add_library(
+  mlx_onnx
+  src/export.cpp
+  src/api.cpp
+  src/compat.cpp
+  src/io.cpp
+  src/lowering.cpp
+  src/mappings.cpp
+  src/onnx.cpp
+  src/shared.cpp)
+
+set_target_properties(mlx_onnx PROPERTIES OUTPUT_NAME mlx_onnx)
+
+target_include_directories(
+  mlx_onnx
+  PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
+         $<INSTALL_INTERFACE:include>
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
+
+if(MLX_ONNX_USE_EXTERNAL_MLX)
+  target_include_directories(mlx_onnx PRIVATE ${MLX_ONNX_EXTERNAL_MLX_INCLUDE_DIR})
+endif()
+
+target_link_libraries(mlx_onnx PUBLIC mlx nlohmann_json::nlohmann_json)
+
+if(MLX_ONNX_BUILD_PYTHON_BINDINGS)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/python/src)
+  set(MLX_ONNX_PY_INIT_FILE ${CMAKE_CURRENT_SOURCE_DIR}/python/mlx_onnx/__init__.py)
+  if(NOT EXISTS ${MLX_ONNX_PY_INIT_FILE})
+    set(MLX_ONNX_PY_INIT_FILE ${CMAKE_CURRENT_BINARY_DIR}/mlx_onnx___init__.py)
+    file(WRITE ${MLX_ONNX_PY_INIT_FILE} "from ._core import *  # noqa: F401,F403\n")
+  endif()
+  if(NOT EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/mlx/CMakeLists.txt)
+    message(FATAL_ERROR "Bundled mlx sources are missing at ${CMAKE_CURRENT_SOURCE_DIR}/mlx")
+  endif()
+  if(NOT TARGET core)
+    message(FATAL_ERROR "Bundled mlx Python extension target `core` was not built")
+  endif()
+  install(TARGETS core LIBRARY DESTINATION mlx COMPONENT python)
+  if(APPLE AND MLX_BUILD_METAL)
+    # MLX looks for mlx.metallib next to the extension module using MLX runtime.
+    install(
+      FILES ${CMAKE_CURRENT_BINARY_DIR}/mlx/mlx/backend/metal/kernels/mlx.metallib
+      DESTINATION mlx
+      COMPONENT python)
+    # mlx_onnx._core also links MLX and resolves the same metallib at runtime.
+    install(
+      FILES ${CMAKE_CURRENT_BINARY_DIR}/mlx/mlx/backend/metal/kernels/mlx.metallib
+      DESTINATION mlx_onnx
+      COMPONENT python)
+  endif()
+  install(
+    DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/mlx/python/mlx/
+    DESTINATION mlx
+    COMPONENT python
+    PATTERN "__pycache__" EXCLUDE)
+  install(
+    FILES ${MLX_ONNX_PY_INIT_FILE}
+    DESTINATION mlx_onnx
+    RENAME __init__.py
+    COMPONENT python)
+  set(MLX_ONNX_VENDOR_MLX_ROOT mlx_onnx/_vendor/mlx)
+  install(
+    FILES ${CMAKE_CURRENT_SOURCE_DIR}/mlx/CMakeLists.txt
+          ${CMAKE_CURRENT_SOURCE_DIR}/mlx/mlx.pc.in
+          ${CMAKE_CURRENT_SOURCE_DIR}/mlx/LICENSE
+          ${CMAKE_CURRENT_SOURCE_DIR}/mlx/ACKNOWLEDGMENTS.md
+    DESTINATION ${MLX_ONNX_VENDOR_MLX_ROOT}
+    COMPONENT python)
+  install(
+    DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/mlx/cmake
+              ${CMAKE_CURRENT_SOURCE_DIR}/mlx/mlx
+    DESTINATION ${MLX_ONNX_VENDOR_MLX_ROOT}
+    COMPONENT python)
+endif()
+
+if(MLX_ONNX_INSTALL_CPP_ARTIFACTS)
+  include(GNUInstallDirs)
+  install(
+    TARGETS mlx_onnx
+    LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
+    INCLUDES DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
+    COMPONENT cpp)
+
+  install(DIRECTORY include/ DESTINATION ${CMAKE_INSTALL_INCLUDEDIR} COMPONENT cpp)
+endif()

data/submodules/mlx-onnx/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 MLX Contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

data/submodules/mlx-onnx/include/mlx/ir.hpp

@@ -0,0 +1,88 @@
+#pragma once
+
+#include <cstdint>
+#include <functional>
+#include <string>
+#include <vector>
+
+#include "mlx/api.h"
+#include "mlx/array.h"
+#include "mlx/export.h"
+
+namespace mlx::onnx {
+
+constexpr int64_t kGraphIrVersion = 1;
+
+using IrCaptureFunction = std::function<std::vector<mlx::core::array>(
+    const mlx::core::Args&,
+    const mlx::core::Kwargs&)>;
+
+struct OnnxBinaryWriteOptions {
+  bool external_data = false;
+  std::string external_data_file = "weights.bin";
+  int64_t external_data_size_threshold = 1024;
+};
+
+struct OnnxBinaryArtifact {
+  std::string model_bytes;
+  std::string external_data_bytes;
+  bool has_external_data = false;
+};
+
+MLX_API std::string ir_to_onnx_json(
+    const std::string& ir_json,
+    int64_t opset,
+    const std::string& model_name);
+
+MLX_API std::string ir_to_onnx(
+    const std::string& target_path,
+    const std::string& ir_json,
+    int64_t opset,
+    const std::string& model_name,
+    const OnnxBinaryWriteOptions& options);
+
+MLX_API std::string ir_compatibility_report_json(
+    const std::string& ir_json);
+
+MLX_API std::string export_ir_json(
+    const IrCaptureFunction& fun,
+    const mlx::core::Args& args,
+    const mlx::core::Kwargs& kwargs,
+    bool shapeless);
+
+MLX_API std::string export_onnx_compatibility_report_json(
+    const IrCaptureFunction& fun,
+    const mlx::core::Args& args,
+    const mlx::core::Kwargs& kwargs,
+    bool shapeless);
+
+MLX_API std::string export_onnx_json(
+    const IrCaptureFunction& fun,
+    const mlx::core::Args& args,
+    const mlx::core::Kwargs& kwargs,
+    bool shapeless,
+    int64_t opset,
+    const std::string& model_name);
+
+MLX_API std::string export_onnx(
+    const std::string& target_path,
+    const IrCaptureFunction& fun,
+    const mlx::core::Args& args,
+    const mlx::core::Kwargs& kwargs,
+    bool shapeless,
+    int64_t opset,
+    const std::string& model_name,
+    const OnnxBinaryWriteOptions& options);
+
+MLX_API OnnxBinaryArtifact build_onnx_binary_artifact_from_onnx_json(
+    const std::string& onnx_json,
+    const OnnxBinaryWriteOptions& options);
+
+MLX_API std::string write_onnx_binary_artifact_to_path(
+    const std::string& target_path,
+    const OnnxBinaryArtifact& artifact,
+    const OnnxBinaryWriteOptions& options);
+
+MLX_API bool ir_is_unsupported_error_message(const std::string& message);
+
+} // namespace mlx::onnx

data/submodules/mlx-onnx/src/api.cpp

@@ -0,0 +1,81 @@
+#include "detail.hpp"
+#include "json.hpp"
+
+#include <sstream>
+#include <stdexcept>
+
+namespace mlx::onnx {
+namespace {
+
+OrderedJson parse_json_payload_from_string(
+    const std::string& raw,
+    const char* label) {
+  try {
+    return OrderedJson::parse(raw);
+  } catch (const std::exception& error) {
+    std::ostringstream out;
+    out << "failed to parse " << label << ": " << error.what();
+    throw std::invalid_argument(
+        detail::tagged_error_message("ir.api", out.str()));
+  }
+}
+
+template <typename Result, typename Callable>
+Result with_ir_api_error_tag(Callable&& callable) {
+  try {
+    return callable();
+  } catch (const std::exception& error) {
+    throw std::runtime_error(
+        detail::tagged_error_message("ir.api", error.what()));
+  }
+}
+
+std::string ir_to_onnx_json_impl(
+    const std::string& ir_json,
+    int64_t opset,
+    const std::string& model_name) {
+  const auto payload =
+      parse_json_payload_from_string(ir_json, "graph ir json");
+  return ir_to_onnx_json_payload(payload, opset, model_name).dump();
+}
+
+std::string ir_compatibility_report_json_impl(
+    const std::string& ir_json) {
+  const auto payload =
+      parse_json_payload_from_string(ir_json, "graph ir json");
+  return ir_compatibility_report_payload(payload).dump();
+}
+
+OnnxBinaryArtifact build_onnx_binary_artifact_from_onnx_json_impl(
+    const std::string& onnx_json,
+    const OnnxBinaryWriteOptions& options) {
+  const auto payload = parse_json_payload_from_string(onnx_json, "onnx json");
+  return build_onnx_binary_artifact_from_stub(payload, options);
+}
+
+} // namespace
+
+std::string ir_to_onnx_json(
+    const std::string& ir_json,
+    int64_t opset,
+    const std::string& model_name) {
+  return with_ir_api_error_tag<std::string>([&]() {
+    return ir_to_onnx_json_impl(ir_json, opset, model_name);
+  });
+}
+
+std::string ir_compatibility_report_json(
+    const std::string& ir_json) {
+  return with_ir_api_error_tag<std::string>(
+      [&]() { return ir_compatibility_report_json_impl(ir_json); });
+}
+
+OnnxBinaryArtifact build_onnx_binary_artifact_from_onnx_json(
+    const std::string& onnx_json,
+    const OnnxBinaryWriteOptions& options) {
+  return with_ir_api_error_tag<OnnxBinaryArtifact>([&]() {
+    return build_onnx_binary_artifact_from_onnx_json_impl(onnx_json, options);
+  });
+}
+
+} // namespace mlx::onnx

data/submodules/mlx-onnx/src/compat.cpp

@@ -0,0 +1,111 @@
+#include "detail.hpp"
+
+#include <exception>
+#include <limits>
+#include <optional>
+#include <set>
+#include <utility>
+
+namespace mlx::onnx::detail {
+
+OrderedJson ir_compatibility_report_payload_impl(
+    const OrderedJson& payload) {
+  // Simulate lowering node-by-node on cloned state. Successful probes commit
+  // inferred state so later nodes see realistic tensor facts.
+  OrderedJson probe_initializers = OrderedJson::array();
+  auto probe_used_tensor_names = collect_payload_tensor_names(payload);
+  auto probe_known_shapes = collect_known_tensor_shapes(payload);
+  auto probe_known_dtypes = collect_known_tensor_dtypes(payload);
+
+  OrderedJson node_support = OrderedJson::array();
+  size_t unsupported_nodes = 0;
+  std::set<std::string> unsupported_ops;
+
+  const auto& source_nodes = payload.at("nodes");
+  for (size_t index = 0; index < source_nodes.size(); ++index) {
+    const auto& node = source_nodes.at(index);
+    const auto op = node.at("op").get<std::string>();
+
+    bool supported = false;
+    std::optional<std::string> mapped;
+
+    try {
+      auto trial_initializers = probe_initializers;
+      auto trial_used_tensor_names = probe_used_tensor_names;
+      auto trial_known_shapes = probe_known_shapes;
+      auto trial_known_dtypes = probe_known_dtypes;
+      LoweringContext trial_lowering{
+          trial_initializers,
+          trial_used_tensor_names,
+          trial_known_shapes,
+          trial_known_dtypes};
+
+      auto lowered = lower_onnx_node_default(node, index, trial_lowering);
+      if (!lowered.empty() && lowered.front().contains("op_type") &&
+          lowered.front().at("op_type").is_string()) {
+        mapped = lowered.front().at("op_type").get<std::string>();
+      } else {
+        mapped = onnx_op_type_for_node(node, false, &trial_known_shapes);
+      }
+
+      probe_initializers = std::move(trial_initializers);
+      probe_used_tensor_names = std::move(trial_used_tensor_names);
+      probe_known_shapes = std::move(trial_known_shapes);
+      probe_known_dtypes = std::move(trial_known_dtypes);
+      supported = true;
+    } catch (const std::exception&) {
+      try {
+        mapped = onnx_op_type_for_node(node, false, &probe_known_shapes);
+      } catch (const std::exception&) {
+        mapped = std::nullopt;
+      }
+    }
+
+    OrderedJson entry = OrderedJson::object();
+    entry["index"] = index;
+    entry["op"] = op;
+    entry["supported"] = supported;
+    if (mapped.has_value()) {
+      entry["onnx_op_type"] = mapped.value();
+    } else {
+      entry["onnx_op_type"] = nullptr;
+    }
+    node_support.push_back(std::move(entry));
+
+    if (!supported) {
+      ++unsupported_nodes;
+      unsupported_ops.insert(op);
+    }
+  }
+
+  int64_t ir_version = kGraphIrVersion;
+  if (payload.contains("ir_version")) {
+    const auto& value = payload.at("ir_version");
+    if (value.is_number_integer()) {
+      ir_version = value.get<int64_t>();
+    } else if (value.is_number_unsigned()) {
+      const auto raw = value.get<uint64_t>();
+      if (raw <= static_cast<uint64_t>(std::numeric_limits<int64_t>::max())) {
+        ir_version = static_cast<int64_t>(raw);
+      }
+    }
+  }
+
+  OrderedJson unsupported_ops_json = OrderedJson::array();
+  for (const auto& unsupported_op : unsupported_ops) {
+    unsupported_ops_json.push_back(unsupported_op);
+  }
+
+  OrderedJson report = OrderedJson::object();
+  report["format"] = "webgpu_compat_report_v1";
+  report["ir_version"] = ir_version;
+  report["total_nodes"] = source_nodes.size();
+  report["supported_nodes"] = source_nodes.size() - unsupported_nodes;
+  report["unsupported_nodes"] = unsupported_nodes;
+  report["unsupported_ops"] = std::move(unsupported_ops_json);
+  report["ready_for_stub_conversion"] = unsupported_nodes == 0;
+  report["nodes"] = std::move(node_support);
+  return report;
+}
+
+} // namespace mlx::onnx::detail

data/submodules/mlx-onnx/src/detail.hpp

@@ -0,0 +1,69 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <map>
+#include <optional>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "json.hpp"
+
+namespace mlx::onnx::detail {
+
+using Shape = std::vector<int64_t>;
+using ShapeMap = std::map<std::string, Shape>;
+using DtypeMap = std::map<std::string, std::string>;
+using NameSet = std::set<std::string>;
+
+std::string tagged_error_message(
+    const std::string& tag,
+    const std::string& message);
+bool json_is_numeric(const OrderedJson& value);
+int64_t normalized_integer_scalar(
+    const OrderedJson& value,
+    const std::string& label);
+std::vector<int64_t> normalize_integer_vector(
+    const OrderedJson& value,
+    const std::string& label);
+std::vector<std::string> parse_string_array(
+    const OrderedJson& value,
+    const std::string& label);
+std::string canonical_dtype(const std::string& dtype);
+std::optional<std::string> canonical_dtype(
+    const std::optional<std::string>& dtype);
+std::optional<std::string> onnx_effective_dtype(
+    const std::optional<std::string>& dtype);
+std::string onnx_effective_dtype(const std::string& dtype);
+std::string onnx_dtype_symbol(const std::string& dtype);
+void for_each_declared_payload_tensor(
+    const OrderedJson& payload,
+    const std::function<void(const OrderedJson&)>& visitor);
+
+struct LoweringContext {
+  OrderedJson& initializers;
+  NameSet& used_tensor_names;
+  ShapeMap& known_shapes;
+  DtypeMap& known_dtypes;
+};
+
+NameSet collect_payload_tensor_names(const OrderedJson& payload);
+ShapeMap collect_known_tensor_shapes(const OrderedJson& payload);
+DtypeMap collect_known_tensor_dtypes(const OrderedJson& payload);
+std::optional<std::string> onnx_op_type_for_node(
+    const OrderedJson& node,
+    bool strict,
+    const ShapeMap* known_shapes);
+std::vector<OrderedJson> lower_onnx_node_default(
+    const OrderedJson& node,
+    size_t node_index,
+    LoweringContext& lowering);
+OrderedJson ir_compatibility_report_payload_impl(
+    const OrderedJson& payload);
+OnnxBinaryArtifact build_onnx_binary_artifact_from_stub_impl(
+    const OrderedJson& onnx_stub,
+    const OnnxBinaryWriteOptions& options);
+
+} // namespace mlx::onnx::detail