mlx 0.30.7.2 → 0.30.7.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 215a912d2353fd5edaa60e320a5b857aa13009e55f3a190b21b2ffe5735f37af
-  data.tar.gz: 1c9b4279f8077e3cd067354ea692e1b24248afe7262d88d4569c566c52f5a158
+  metadata.gz: d8190b13b012fe0693ef46cc3f20b01b78f6d13cde44674fd01999434e56eda9
+  data.tar.gz: '059a7b993b17cae7bd448567ef6a9058d0a575cdc143a3784b411953e1309e7d'
 SHA512:
-  metadata.gz: 66abcbd58ccfc04186df11b0d2b6445c7d1e0ab4a36451742755f6fcf41022363403536c3d27640935364c58231c4c4a03e39ceba97617a59f2ad69acf23dc16
-  data.tar.gz: ba7ad07ccd31e94bdf3fdee73117f069c6ee22c11cdfc3f2470eedee9ee0bc9e976970bab21f156c003b658723a0e3fb2f63f62d0e11d33de3a61fc1d9121711
+  metadata.gz: 1a90a276ac6b68760bdbe348931e6cfc48ec01662a7870522cf344748a8e5df167f3d4f6b9d28bd768abc3c5dc9a5e020abe2a5fc81edf1ef4056e54abc269b7
+  data.tar.gz: 8b17e8c878f32659cc90b365943cf7fce132c03fd7680959c03b615b28819d6a002e6e05236209b0282fdec7b03f06872556f7ac292ac85e00d531fa0bd1c11a

data/ext/mlx/extconf.rb

@@ -3,6 +3,9 @@
 require "etc"
 require "fileutils"
 require "mkmf"
+require "open3"
+require "pathname"
+require "rbconfig"
 
 def run_or_abort(*cmd, chdir:)
   puts ">> #{cmd.join(' ')}"
@@ -17,6 +20,76 @@ def run_with_status(*cmd, chdir:)
   system(*cmd, chdir: chdir)
 end
 
+def git_revision(path)
+  return nil unless Dir.exist?(path)
+
+  stdout, status = Open3.capture2("git", "-C", path, "rev-parse", "HEAD")
+  return nil unless status.success?
+
+  rev = stdout.strip
+  return rev if rev.match?(/\A[0-9a-f]{40}\z/)
+
+  nil
+end
+
+def mlx_revision_pinned_by_mlx_onnx(mlx_onnx_root)
+  return nil unless Dir.exist?(mlx_onnx_root)
+
+  stdout, status = Open3.capture2("git", "-C", mlx_onnx_root, "submodule", "status", "--", "mlx")
+  return nil unless status.success?
+
+  line = stdout.lines.first.to_s.strip
+  return nil if line.empty?
+
+  token = line.split(/\s+/).first.to_s
+  token = token.delete_prefix("-").delete_prefix("+").delete_prefix("U")
+  return token if token.match?(/\A[0-9a-f]{40}\z/)
+
+  nil
+end
+
+def enforce_mlx_onnx_compatibility!(mlx_root:, mlx_onnx_root:)
+  workspace_mlx_revision =
+    ENV.fetch("MLX_EXTCONF_TEST_MLX_REVISION", "").strip
+  pinned_mlx_revision =
+    ENV.fetch("MLX_EXTCONF_TEST_MLX_ONNX_PINNED_MLX_REVISION", "").strip
+
+  workspace_mlx_revision = git_revision(mlx_root) if workspace_mlx_revision.empty?
+  pinned_mlx_revision = mlx_revision_pinned_by_mlx_onnx(mlx_onnx_root) if pinned_mlx_revision.empty?
+
+  return if workspace_mlx_revision.nil? || pinned_mlx_revision.nil?
+  return if workspace_mlx_revision == pinned_mlx_revision
+
+  abort(<<~MSG)
+    mlx/mlx-onnx revision mismatch detected.
+    workspace mlx revision: #{workspace_mlx_revision}
+    mlx-onnx pinned mlx revision: #{pinned_mlx_revision}
+    Run:
+      git submodule update --init --recursive submodules/mlx submodules/mlx-onnx
+  MSG
+end
+
+def patch_mlx_onnx_gcc_optional_shape_initlist!(mlx_onnx_root)
+  lowering_cpp = File.join(mlx_onnx_root, "src", "lowering.cpp")
+  return unless File.file?(lowering_cpp)
+
+  source = File.read(lowering_cpp)
+  patched = source
+    .gsub(
+      "std::optional<Shape>({work_shape[0], 1, seq_len})",
+      "std::optional<Shape>(Shape{work_shape[0], 1, seq_len})"
+    )
+    .gsub(
+      "std::optional<Shape>({seq_len})",
+      "std::optional<Shape>(Shape{seq_len})"
+    )
+
+  return if patched == source
+
+  File.write(lowering_cpp, patched)
+  puts "patched mlx-onnx lowering.cpp optional Shape initlists for GCC compatibility"
+end
+
 def rpath_flag(path)
   case RUBY_PLATFORM
   when /darwin/
@@ -28,15 +101,120 @@ def rpath_flag(path)
   end
 end
 
+def cmake_compilers_from_cache(cache_path)
+  return [nil, nil] unless File.file?(cache_path)
+
+  cc = nil
+  cxx = nil
+  File.foreach(cache_path) do |line|
+    cc = Regexp.last_match(1).strip if line =~ /^CMAKE_C_COMPILER:(?:STRING|FILEPATH)=(.+)$/
+    cxx = Regexp.last_match(1).strip if line =~ /^CMAKE_CXX_COMPILER:(?:STRING|FILEPATH)=(.+)$/
+  end
+  [cc, cxx]
+end
+
+def normalize_compiler_for_mkmf(path, kind)
+  return path if path.nil? || path.empty?
+
+  case kind
+  when :cc
+    return "/usr/bin/clang" if path.end_with?("/usr/bin/cc")
+  when :cxx
+    return "/usr/bin/clang++" if path.end_with?("/usr/bin/c++")
+  end
+  path
+end
+
+def force_mkmf_compilers!(cc, cxx)
+  return if cc.nil? || cc.empty? || cxx.nil? || cxx.empty?
+
+  cc = normalize_compiler_for_mkmf(cc, :cc)
+  cxx = normalize_compiler_for_mkmf(cxx, :cxx)
+
+  [RbConfig::CONFIG, RbConfig::MAKEFILE_CONFIG].each do |cfg|
+    cfg["CC"] = cc if cfg.key?("CC")
+    cfg["CXX"] = cxx if cfg.key?("CXX")
+
+    if cfg.key?("LDSHARED") && cfg["LDSHARED"]
+      cfg["LDSHARED"] = cfg["LDSHARED"].sub(/\A\S+/, cc)
+    end
+    if cfg.key?("LDSHAREDXX") && cfg["LDSHAREDXX"]
+      cfg["LDSHAREDXX"] = cfg["LDSHAREDXX"].sub(/\A\S+/, cxx)
+    end
+  end
+
+  $CC = cc if defined?($CC)
+  $CXX = cxx if defined?($CXX)
+  $LDSHARED = $LDSHARED.sub(/\A\S+/, cc) if defined?($LDSHARED) && $LDSHARED
+  $LDSHAREDXX = $LDSHAREDXX.sub(/\A\S+/, cxx) if defined?($LDSHAREDXX) && $LDSHAREDXX
+end
+
+def patch_makefile_compilers!(makefile_path, cc, cxx)
+  return if cc.nil? || cc.empty? || cxx.nil? || cxx.empty?
+  return unless File.file?(makefile_path)
+
+  cc = normalize_compiler_for_mkmf(cc, :cc)
+  cxx = normalize_compiler_for_mkmf(cxx, :cxx)
+
+  text = File.read(makefile_path)
+  text = text.gsub(/^CC = .+$/, "CC = #{cc}")
+  text = text.gsub(/^CXX = .+$/, "CXX = #{cxx}")
+  File.write(makefile_path, text)
+end
+
+def patch_makefile_sources!(makefile_path, source_files, header_files)
+  return unless File.file?(makefile_path)
+
+  objects = source_files.map do |path|
+    stem = path.sub(/\.[^.]+\z/, "")
+    "#{stem}.o"
+  end
+  text = File.read(makefile_path)
+  text = text.gsub(/^ORIG_SRCS = .+$/, "ORIG_SRCS = #{source_files.join(' ')}")
+  text = text.gsub(/^OBJS = .+$/, "OBJS = #{objects.join(' ')}")
+  text = text.gsub(/^HDRS = .+$/, "HDRS = #{header_files.map { |path| '$(srcdir)/' + path }.join(' ')}")
+  File.write(makefile_path, text)
+end
+
+def patch_makefile_include_dirs!(makefile_path, include_dirs)
+  return unless File.file?(makefile_path)
+
+  text = File.read(makefile_path)
+  match = text.match(/^CPPFLAGS = (.+)$/)
+  return if match.nil?
+
+  cppflags = match[1]
+  include_dirs.each do |path|
+    flag = "-I#{path}"
+    cppflags = "#{cppflags} #{flag}" unless cppflags.include?(flag)
+  end
+  text = text.sub(/^CPPFLAGS = .+$/, "CPPFLAGS = #{cppflags}")
+  File.write(makefile_path, text)
+end
+
 repo_root = File.expand_path("../..", __dir__)
-mlx_root = File.join(repo_root, "mlx")
-mlx_include_dir = mlx_root
+submodules_root = File.join(repo_root, "submodules")
+mlx_root = File.join(submodules_root, "mlx")
+mlx_onnx_root = File.join(submodules_root, "mlx-onnx")
+mlx_public_include_dir = mlx_root
+mlx_internal_include_dir = File.join(mlx_root, "mlx")
+mlx_onnx_public_include_dir = File.join(mlx_onnx_root, "include")
+mlx_onnx_internal_include_dir = File.join(mlx_onnx_root, "src")
 ext_root = File.expand_path(__dir__)
+ext_onnx_root = File.join(repo_root, "ext", "mlx-onnx")
 build_root = File.join(ext_root, "build")
 mlx_build_dir = File.join(build_root, "mlx")
+mlx_onnx_build_dir = File.join(build_root, "mlx-onnx")
 mlx_install_dir = File.join(build_root, "install")
 jobs = [Etc.nprocessors, 1].max
 
+enforce_mlx_onnx_compatibility!(mlx_root: mlx_root, mlx_onnx_root: mlx_onnx_root)
+patch_mlx_onnx_gcc_optional_shape_initlist!(mlx_onnx_root)
+if ENV["MLX_EXTCONF_VALIDATE_ONLY"] == "1"
+  puts "mlx-onnx compatibility check passed"
+  exit 0
+end
+
 FileUtils.mkdir_p(mlx_build_dir)
 
 cmake_configure = [
@@ -54,7 +232,7 @@ cmake_configure = [
   "-DMLX_BUILD_PYTHON_STUBS=OFF",
   "-DMLX_BUILD_METAL=ON",
   "-DMLX_BUILD_GGUF=OFF",
-  "-DMLX_BUILD_SAFETENSORS=OFF",
+  "-DMLX_BUILD_SAFETENSORS=ON",
   "-DBUILD_SHARED_LIBS=ON"
 ]
 
@@ -78,17 +256,98 @@ unless configured
 end
 run_or_abort(*cmake_build, chdir: ext_root)
 
-include_dir = mlx_include_dir
+cmake_cache_path = File.join(mlx_build_dir, "CMakeCache.txt")
+cmake_cc, cmake_cxx = cmake_compilers_from_cache(cmake_cache_path)
+force_mkmf_compilers!(cmake_cc, cmake_cxx)
+
+abort("missing MLX include dir: #{mlx_public_include_dir}") unless Dir.exist?(mlx_public_include_dir)
+abort("missing MLX internal include dir: #{mlx_internal_include_dir}") unless Dir.exist?(mlx_internal_include_dir)
+abort("missing mlx-onnx include dir: #{mlx_onnx_public_include_dir}") unless Dir.exist?(mlx_onnx_public_include_dir)
+abort("missing mlx-onnx internal include dir: #{mlx_onnx_internal_include_dir}") unless Dir.exist?(mlx_onnx_internal_include_dir)
+include_dirs = [
+  mlx_public_include_dir,
+  mlx_internal_include_dir,
+  mlx_onnx_public_include_dir,
+  mlx_onnx_internal_include_dir
+]
 lib_dir = File.join(mlx_install_dir, "lib")
 
-abort("missing MLX include dir: #{include_dir}") unless Dir.exist?(include_dir)
+cmake_onnx_configure = [
+  "cmake",
+  "-S",
+  mlx_onnx_root,
+  "-B",
+  mlx_onnx_build_dir,
+  "-DCMAKE_BUILD_TYPE=Release",
+  "-DCMAKE_INSTALL_PREFIX=#{mlx_install_dir}",
+  "-DMLX_ONNX_USE_EXTERNAL_MLX=ON",
+  "-DMLX_ONNX_EXTERNAL_MLX_INCLUDE_DIR=#{mlx_public_include_dir}",
+  "-DMLX_ONNX_EXTERNAL_MLX_LIB_DIR=#{lib_dir}",
+  "-DMLX_ONNX_BUILD_PYTHON_BINDINGS=OFF"
+]
+
+cmake_onnx_build = [
+  "cmake",
+  "--build",
+  mlx_onnx_build_dir,
+  "--target",
+  "install",
+  "--config",
+  "Release",
+  "-j#{jobs}"
+]
+
+run_or_abort(*cmake_onnx_configure, chdir: ext_root)
+run_or_abort(*cmake_onnx_build, chdir: ext_root)
+
+json_include_dirs = [
+  File.join(mlx_build_dir, "_deps", "json-src", "include"),
+  File.join(mlx_build_dir, "_deps", "json-src", "single_include"),
+  File.join(mlx_build_dir, "_deps", "json-src", "single_include", "nlohmann"),
+  File.join(mlx_onnx_build_dir, "_deps", "nlohmann_json-src", "include")
+].select { |path| Dir.exist?(path) }
+
 abort("missing MLX lib dir: #{lib_dir}") unless Dir.exist?(lib_dir)
 
-dir_config("mlx", include_dir, lib_dir)
+dir_config("mlx", include_dirs.first, lib_dir)
 
 $CXXFLAGS = "#{$CXXFLAGS} -std=c++20"
-$CPPFLAGS = "#{$CPPFLAGS} -I#{include_dir}"
+include_dirs.each do |path|
+  $CPPFLAGS = "#{$CPPFLAGS} -I#{path}"
+end
+json_include_dirs.each do |path|
+  $CPPFLAGS = "#{$CPPFLAGS} -I#{path}"
+end
 $LDFLAGS = "#{$LDFLAGS} -L#{lib_dir} #{rpath_flag(lib_dir)}"
-$libs = "-lmlx #{$libs}"
+$libs = "-lmlx_onnx -lmlx #{$libs}"
+
+source_files = (
+  Dir.glob(File.join(ext_root, "**", "*.{c,cc,cpp,cxx}")) +
+  Dir.glob(File.join(ext_onnx_root, "**", "*.{c,cc,cpp,cxx}"))
+).reject { |path| path.start_with?(File.join(build_root, "")) }
+ .map { |path| Pathname(path).relative_path_from(Pathname(ext_root)).to_s }
+ .sort
+header_files = (
+  Dir.glob(File.join(ext_root, "**", "*.{h,hpp,hh}")) +
+  Dir.glob(File.join(ext_onnx_root, "**", "*.{h,hpp,hh}"))
+).reject { |path| path.start_with?(File.join(build_root, "")) }
+ .map { |path| Pathname(path).relative_path_from(Pathname(ext_root)).to_s }
+ .sort
+$srcs = source_files
+$objs = source_files.map { |path| "#{path.sub(/\.[^.]+\z/, "")}.o" }
+$hdrs = header_files
 
 create_makefile("mlx/native")
+makefile_path = File.join(ext_root, "Makefile")
+patch_makefile_compilers!(makefile_path, cmake_cc, cmake_cxx)
+patch_makefile_sources!(makefile_path, source_files, header_files)
+patch_makefile_include_dirs!(
+  makefile_path,
+  [
+    mlx_internal_include_dir,
+    mlx_onnx_public_include_dir,
+    mlx_onnx_internal_include_dir,
+    File.join(mlx_build_dir, "_deps", "json-src", "single_include", "nlohmann"),
+    File.join(mlx_onnx_build_dir, "_deps", "nlohmann_json-src", "include")
+  ]
+)

data/ext/mlx/native.cpp

@@ -10,12 +10,15 @@
 #include <limits>
 #include <optional>
 #include <sstream>
+#include <stdexcept>
 #include <string>
 #include <type_traits>
 #include <unordered_map>
 #include <variant>
 #include <vector>
 
+#include <nlohmann/json.hpp>
+
 #include "mlx/array.h"
 #include "mlx/backend/metal/metal.h"
 #include "mlx/compile.h"
@@ -37,6 +40,7 @@
 #include "mlx/transforms.h"
 #include "mlx/utils.h"
 #include "mlx/version.h"
+#include "../mlx-onnx/native.hpp"
 
 namespace mx = mlx::core;
 namespace mxfft = mlx::core::fft;
@@ -44,6 +48,7 @@ namespace mxfast = mlx::core::fast;
 namespace mxlinalg = mlx::core::linalg;
 namespace mxmetal = mlx::core::metal;
 namespace mxdist = mlx::core::distributed;
+using OrderedJson = nlohmann::ordered_json;
 
 static VALUE mMLX;
 static VALUE mNative;
@@ -339,14 +344,7 @@ static VALUE id_to_symbol(ID id) {
 }
 
 static ID cached_intern_id(const char* name) {
-  static std::unordered_map<std::string, ID> cache;
-  auto it = cache.find(name);
-  if (it != cache.end()) {
-    return it->second;
-  }
-  const ID id = rb_intern(name);
-  cache.emplace(name, id);
-  return id;
+  return rb_intern(name);
 }
 
 static mx::Device::DeviceType device_type_from_value(VALUE value) {
@@ -928,12 +926,6 @@ struct ArrayCollector {
 
 static void collect_arrays_from_tree(VALUE value, std::vector<mx::array>& arrays);
 
-static int hash_collect_arrays_iter(VALUE, VALUE value, VALUE arg) {
-  auto* collector = reinterpret_cast<ArrayCollector*>(arg);
-  collect_arrays_from_tree(value, *collector->arrays);
-  return ST_CONTINUE;
-}
-
 static void collect_arrays_from_tree(VALUE value, std::vector<mx::array>& arrays) {
   if (rb_obj_is_kind_of(value, cArray)) {
     arrays.push_back(array_unwrap(value));
@@ -947,28 +939,37 @@ static void collect_arrays_from_tree(VALUE value, std::vector<mx::array>& arrays
     return;
   }
   if (RB_TYPE_P(value, T_HASH)) {
-    ArrayCollector collector{&arrays};
-    rb_hash_foreach(value, hash_collect_arrays_iter, reinterpret_cast<VALUE>(&collector));
+    VALUE keys = rb_funcall(value, cached_intern_id("keys"), 0);
+    const long len = RARRAY_LEN(keys);
+    for (long i = 0; i < len; ++i) {
+      VALUE key = rb_ary_entry(keys, i);
+      VALUE hash_value = rb_hash_lookup2(value, key, Qundef);
+      if (hash_value == Qundef) {
+        continue;
+      }
+      collect_arrays_from_tree(hash_value, arrays);
+    }
   }
 }
 
-struct ArrayMapBuilder {
-  std::unordered_map<std::string, mx::array> map;
-};
-
-static int hash_to_array_map_iter(VALUE key, VALUE value, VALUE arg) {
-  auto* builder = reinterpret_cast<ArrayMapBuilder*>(arg);
-  builder->map.insert_or_assign(string_from_ruby(key), array_unwrap(value));
-  return ST_CONTINUE;
-}
-
 static std::unordered_map<std::string, mx::array> array_map_from_ruby_hash(VALUE value) {
   if (!RB_TYPE_P(value, T_HASH)) {
     rb_raise(rb_eTypeError, "expected Hash mapping String/Symbol keys to MLX::Core::Array");
   }
-  ArrayMapBuilder builder;
-  rb_hash_foreach(value, hash_to_array_map_iter, reinterpret_cast<VALUE>(&builder));
-  return builder.map;
+  VALUE keys = rb_funcall(value, cached_intern_id("keys"), 0);
+  const long len = RARRAY_LEN(keys);
+  std::unordered_map<std::string, mx::array> out;
+  out.reserve(static_cast<size_t>(len));
+  for (long i = 0; i < len; ++i) {
+    VALUE key = rb_ary_entry(keys, i);
+    VALUE hash_value = rb_hash_lookup2(value, key, Qundef);
+    if (hash_value == Qundef) {
+      continue;
+    }
+    std::string ruby_key = string_from_ruby(key);
+    out.insert_or_assign(ruby_key, array_unwrap(hash_value));
+  }
+  return out;
 }
 
 static VALUE ruby_hash_of_arrays(const std::unordered_map<std::string, mx::array>& map) {
@@ -990,16 +991,6 @@ static VALUE ruby_hash_of_strings(const std::unordered_map<std::string, std::str
   return out;
 }
 
-struct StringMapBuilder {
-  std::unordered_map<std::string, std::string> map;
-};
-
-static int hash_to_string_map_iter(VALUE key, VALUE value, VALUE arg) {
-  auto* builder = reinterpret_cast<StringMapBuilder*>(arg);
-  builder->map.insert_or_assign(string_from_ruby(key), string_from_ruby(value));
-  return ST_CONTINUE;
-}
-
 static std::unordered_map<std::string, std::string> string_map_from_ruby_hash(VALUE value) {
   if (NIL_P(value)) {
     return {};
@@ -1007,9 +998,20 @@ static std::unordered_map<std::string, std::string> string_map_from_ruby_hash(VA
   if (!RB_TYPE_P(value, T_HASH)) {
     rb_raise(rb_eTypeError, "expected Hash mapping String/Symbol keys to String values");
   }
-  StringMapBuilder builder;
-  rb_hash_foreach(value, hash_to_string_map_iter, reinterpret_cast<VALUE>(&builder));
-  return builder.map;
+  VALUE keys = rb_funcall(value, cached_intern_id("keys"), 0);
+  const long len = RARRAY_LEN(keys);
+  std::unordered_map<std::string, std::string> out;
+  out.reserve(static_cast<size_t>(len));
+  for (long i = 0; i < len; ++i) {
+    VALUE key = rb_ary_entry(keys, i);
+    VALUE hash_value = rb_hash_lookup2(value, key, Qundef);
+    if (hash_value == Qundef) {
+      continue;
+    }
+    std::string ruby_key = string_from_ruby(key);
+    out.insert_or_assign(ruby_key, string_from_ruby(hash_value));
+  }
+  return out;
 }
 
 static mx::GGUFMetaData gguf_metadata_from_ruby(VALUE value) {
@@ -1038,16 +1040,6 @@ static mx::GGUFMetaData gguf_metadata_from_ruby(VALUE value) {
   return std::monostate{};
 }
 
-struct GGUFMetaMapBuilder {
-  std::unordered_map<std::string, mx::GGUFMetaData> map;
-};
-
-static int hash_to_gguf_meta_map_iter(VALUE key, VALUE value, VALUE arg) {
-  auto* builder = reinterpret_cast<GGUFMetaMapBuilder*>(arg);
-  builder->map.insert_or_assign(string_from_ruby(key), gguf_metadata_from_ruby(value));
-  return ST_CONTINUE;
-}
-
 static std::unordered_map<std::string, mx::GGUFMetaData> gguf_meta_map_from_ruby_hash(VALUE value) {
   if (NIL_P(value)) {
     return {};
@@ -1055,9 +1047,20 @@ static std::unordered_map<std::string, mx::GGUFMetaData> gguf_meta_map_from_ruby
   if (!RB_TYPE_P(value, T_HASH)) {
     rb_raise(rb_eTypeError, "expected Hash for GGUF metadata");
   }
-  GGUFMetaMapBuilder builder;
-  rb_hash_foreach(value, hash_to_gguf_meta_map_iter, reinterpret_cast<VALUE>(&builder));
-  return builder.map;
+  VALUE keys = rb_funcall(value, cached_intern_id("keys"), 0);
+  const long len = RARRAY_LEN(keys);
+  std::unordered_map<std::string, mx::GGUFMetaData> out;
+  out.reserve(static_cast<size_t>(len));
+  for (long i = 0; i < len; ++i) {
+    VALUE key = rb_ary_entry(keys, i);
+    VALUE hash_value = rb_hash_lookup2(value, key, Qundef);
+    if (hash_value == Qundef) {
+      continue;
+    }
+    std::string ruby_key = string_from_ruby(key);
+    out.insert_or_assign(ruby_key, gguf_metadata_from_ruby(hash_value));
+  }
+  return out;
 }
 
 static VALUE gguf_metadata_to_ruby(const mx::GGUFMetaData& value) {
@@ -1501,6 +1504,23 @@ args_kwargs_function_from_callable(VALUE callable) {
   };
 }
 
+mx::array onnx_array_from_ruby(VALUE value) {
+  return array_from_ruby(value, std::nullopt);
+}
+
+std::vector<mx::array> onnx_array_vector_from_ruby(VALUE value) {
+  return array_vector_from_ruby(value);
+}
+
+std::unordered_map<std::string, mx::array> onnx_array_map_from_ruby_hash(VALUE value) {
+  return array_map_from_ruby_hash(value);
+}
+
+std::function<std::vector<mx::array>(const mx::Args&, const mx::Kwargs&)>
+onnx_args_kwargs_function_from_callable(VALUE callable) {
+  return args_kwargs_function_from_callable(callable);
+}
+
 static std::vector<int> argnums_from_value(VALUE value) {
   if (NIL_P(value)) {
     return {0};
@@ -3495,7 +3515,7 @@ static VALUE core_dequantize(int argc, VALUE* argv, VALUE) {
       dtype = optional_dtype_from_value(argv[6]);
     }
 
-    return array_wrap(mx::dequantize(w, scales, biases, group_size, bits, mode, dtype));
+    return array_wrap(mx::dequantize(w, scales, biases, group_size, bits, mode, std::nullopt, dtype));
   } catch (const std::exception& error) {
     raise_std_exception(error);
     return Qnil;
@@ -6459,6 +6479,30 @@ static VALUE core_identity(int argc, VALUE* argv, VALUE) {
   }
 }
 
+static VALUE core_hanning(int argc, VALUE* argv, VALUE) {
+  try {
+    VALUE m;
+    VALUE stream;
+    rb_scan_args(argc, argv, "11", &m, &stream);
+    return array_wrap(mx::hanning(NUM2INT(m), stream_or_device_from_value(stream)));
+  } catch (const std::exception& error) {
+    raise_std_exception(error);
+    return Qnil;
+  }
+}
+
+static VALUE core_hamming(int argc, VALUE* argv, VALUE) {
+  try {
+    VALUE m;
+    VALUE stream;
+    rb_scan_args(argc, argv, "11", &m, &stream);
+    return array_wrap(mx::hamming(NUM2INT(m), stream_or_device_from_value(stream)));
+  } catch (const std::exception& error) {
+    raise_std_exception(error);
+    return Qnil;
+  }
+}
+
 static VALUE core_tri(int argc, VALUE* argv, VALUE) {
   try {
     if (argc < 1 || argc > 4) {
@@ -6625,7 +6669,8 @@ static VALUE core_clear_cache(VALUE) {
 
 static VALUE core_metal_is_available(VALUE) {
   try {
-    return mxmetal::is_available() ? Qtrue : Qfalse;
+    const mx::Device gpu_device(mx::Device::gpu, 0);
+    return mx::is_available(gpu_device) ? Qtrue : Qfalse;
   } catch (const std::exception& error) {
     raise_std_exception(error);
     return Qnil;
@@ -6654,7 +6699,12 @@ static VALUE core_metal_stop_capture(VALUE) {
 
 static VALUE core_metal_device_info(VALUE) {
   try {
-    const auto& info = mxmetal::device_info();
+    const mx::Device gpu_device(mx::Device::gpu, 0);
+    if (!mx::is_available(gpu_device)) {
+      rb_raise(rb_eRuntimeError, "[metal_device_info] Metal GPU device is not available");
+    }
+
+    const auto& info = mx::device_info(gpu_device);
     VALUE hash = rb_hash_new();
     for (const auto& [key, value] : info) {
       VALUE ruby_key = rb_utf8_str_new(key.c_str(), static_cast<long>(key.size()));
@@ -7616,6 +7666,7 @@ extern "C" void Init_native(void) {
   rb_define_singleton_method(mNative, "loaded?", RUBY_METHOD_FUNC(native_loaded_p), 0);
 
   mCore = rb_define_module_under(mMLX, "Core");
+  init_onnx_native_bindings(mMLX);
   rb_define_singleton_method(mCore, "version", RUBY_METHOD_FUNC(core_version), 0);
 
   rb_define_singleton_method(mCore, "get_active_memory", RUBY_METHOD_FUNC(core_get_active_memory), 0);
@@ -7893,6 +7944,8 @@ extern "C" void Init_native(void) {
   rb_define_singleton_method(mCore, "ones_like", RUBY_METHOD_FUNC(core_ones_like), 1);
   rb_define_singleton_method(mCore, "eye", RUBY_METHOD_FUNC(core_eye), -1);
   rb_define_singleton_method(mCore, "identity", RUBY_METHOD_FUNC(core_identity), -1);
+  rb_define_singleton_method(mCore, "hanning", RUBY_METHOD_FUNC(core_hanning), -1);
+  rb_define_singleton_method(mCore, "hamming", RUBY_METHOD_FUNC(core_hamming), -1);
   rb_define_singleton_method(mCore, "tri", RUBY_METHOD_FUNC(core_tri), -1);
   rb_define_singleton_method(mCore, "tril", RUBY_METHOD_FUNC(core_tril), -1);
   rb_define_singleton_method(mCore, "triu", RUBY_METHOD_FUNC(core_triu), -1);
@@ -8020,4 +8073,5 @@ extern "C" void Init_native(void) {
       "precompiled_cuda_kernel",
       RUBY_METHOD_FUNC(core_precompiled_cuda_kernel),
       -1);
+
 }