mlx 0.30.7.3 → 0.30.7.6

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