mlx 0.30.7.2 → 0.30.7.6

data/ext/mlx-onnx/native.hpp

@@ -0,0 +1,19 @@
+#pragma once
+
+#include <functional>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include <ruby.h>
+
+#include "mlx/array.h"
+#include "mlx/export.h"
+
+mlx::core::array onnx_array_from_ruby(VALUE value);
+std::vector<mlx::core::array> onnx_array_vector_from_ruby(VALUE value);
+std::unordered_map<std::string, mlx::core::array> onnx_array_map_from_ruby_hash(VALUE value);
+std::function<std::vector<mlx::core::array>(const mlx::core::Args&, const mlx::core::Kwargs&)>
+onnx_args_kwargs_function_from_callable(VALUE callable);
+
+extern "C" void init_onnx_native_bindings(VALUE mMLX);

data/lib/mlx/core.rb

@@ -2,6 +2,7 @@
 
 require "open3"
 require "tmpdir"
+require_relative "onnx"
 
 module MLX
   module Core
@@ -334,6 +335,8 @@ module MLX
       alias_method :native_vmap, :vmap if method_defined?(:vmap) && !method_defined?(:native_vmap)
       alias_method :native_export_to_dot,
                    :export_to_dot if method_defined?(:export_to_dot) && !method_defined?(:native_export_to_dot)
+      alias_method :native_array, :array if method_defined?(:array) && !method_defined?(:native_array)
+      alias_method :native_mean, :mean if method_defined?(:mean) && !method_defined?(:native_mean)
 
       %i[savez savez_compressed].each do |method_name|
         if method_defined?(method_name) && instance_method(method_name).owner == self
@@ -343,6 +346,24 @@ module MLX
 
       ARRAY_LEAF = :__mlx_array_leaf__
 
+      def array(value, positional_dtype = nil, dtype: nil)
+        ensure_native!
+        target_dtype = resolve_array_dtype(positional_dtype, dtype)
+        native_array(value, target_dtype)
+      end
+
+      def mean(array, axis = nil, positional_keepdims = nil, keepdims: nil)
+        ensure_native!
+        keepdims_v = resolve_keepdims_argument(positional_keepdims, keepdims)
+        reduced = reduce_mean(array, axis)
+        return reduced unless keepdims_v
+
+        normalize_reduction_axes(array, axis).each do |axis_index|
+          reduced = expand_dims(reduced, axis_index)
+        end
+        reduced
+      end
+
       def load(file, format = nil, return_metadata = false)
         ensure_native!
         format_name = (format || infer_format(file)).to_s
@@ -560,6 +581,97 @@ module MLX
 
       private
 
+      def resolve_array_dtype(positional_dtype, keyword_dtype)
+        normalized_positional = normalize_dtype_alias(positional_dtype)
+        normalized_keyword = normalize_dtype_alias(keyword_dtype)
+        return normalized_keyword if normalized_positional.nil?
+        return normalized_positional if normalized_keyword.nil?
+
+        if dtype_name_for_compare(normalized_positional) != dtype_name_for_compare(normalized_keyword)
+          raise ArgumentError,
+                "array received conflicting dtype arguments (positional=#{positional_dtype.inspect}, keyword=#{keyword_dtype.inspect})"
+        end
+
+        normalized_positional
+      end
+
+      def normalize_dtype_alias(dtype)
+        return nil if dtype.nil?
+        return dtype if dtype.respond_to?(:name)
+        return dtype unless dtype.is_a?(::Symbol) || dtype.is_a?(::String)
+
+        case dtype.to_s.strip.downcase
+        when "bool", "bool_"
+          :bool_
+        when "f16", "fp16", "float16"
+          :float16
+        when "bf16", "bfloat16"
+          :bfloat16
+        when "f32", "fp32", "float32"
+          :float32
+        when "f64", "fp64", "float64"
+          :float64
+        when "c64", "complex64"
+          :complex64
+        else
+          dtype
+        end
+      end
+
+      def dtype_name_for_compare(dtype)
+        return nil if dtype.nil?
+        dtype = normalize_dtype_alias(dtype)
+
+        if dtype.respond_to?(:name)
+          dtype.name.to_s
+        else
+          dtype.to_s
+        end
+      end
+
+      def resolve_keepdims_argument(positional_keepdims, keyword_keepdims)
+        if !positional_keepdims.nil? && !keyword_keepdims.nil? && !!positional_keepdims != !!keyword_keepdims
+          raise ArgumentError,
+                "mean received conflicting keepdims arguments (positional=#{positional_keepdims.inspect}, keyword=#{keyword_keepdims.inspect})"
+        end
+        return !!keyword_keepdims unless keyword_keepdims.nil?
+        return !!positional_keepdims unless positional_keepdims.nil?
+
+        false
+      end
+
+      def reduce_mean(array, axis)
+        if axis.is_a?(::Array)
+          normalize_reduction_axes(array, axis).reverse_each.reduce(array) do |acc, axis_index|
+            native_mean(acc, axis_index)
+          end
+        else
+          native_mean(array, axis)
+        end
+      end
+
+      def normalize_reduction_axes(array, axis)
+        ndim = array.ndim
+        return (0...ndim).to_a if axis.nil?
+
+        raw_axes = axis.is_a?(::Array) ? axis : [axis]
+        axes = raw_axes.map { |entry| normalize_axis_index(entry, ndim) }.sort
+        raise ArgumentError, "axis contains duplicate values: #{raw_axes.inspect}" if axes.uniq.length != axes.length
+
+        axes
+      end
+
+      def normalize_axis_index(axis, ndim)
+        raise TypeError, "axis entries must be Integer" unless axis.is_a?(::Integer)
+
+        out = axis
+        out += ndim if out.negative?
+        if out.negative? || out >= ndim
+          raise ArgumentError, "axis #{axis} is out of bounds for array of dimension #{ndim}"
+        end
+        out
+      end
+
       def infer_format(file)
         path = file_path(file)
         ext = File.extname(path).delete_prefix(".")
@@ -601,7 +713,11 @@ module MLX
           Dir.glob(File.join(dir, "**", "*.npy")).sort.each do |npy_path|
             rel = npy_path.delete_prefix(dir + File::SEPARATOR)
             key = rel.end_with?(".npy") ? rel[0...-4] : rel
-            out[key] = native_load(npy_path, "npy", false)
+            # Force a materialized copy to avoid keeping many file-backed mmap handles open.
+            value = native_load(npy_path, "npy", false)
+            value = add(value, 0)
+            eval(value)
+            out[key] = value
           end
           out
         end
@@ -677,27 +793,95 @@ module MLX
       end
 
       def build_grad_like_function(fun, argnums, argnames, with_value)
+        cache = {}
+
         lambda do |*args, **kwargs|
           selections, flat_inputs = build_target_selections(args, kwargs, argnums, argnames)
-          native_argnums = (0...flat_inputs.length).to_a
-          captured_value = nil
-          lifted = lambda do |*flat_vars|
-            call_args, call_kwargs = apply_flat_vars_to_targets(args, kwargs, selections, flat_vars)
-            raw_value = fun.call(*call_args, **call_kwargs)
-            captured_value = raw_value
-            extract_loss(raw_value)
+          cache_key = grad_selection_cache_key(selections)
+          entry = cache[cache_key]
+          unless entry
+            call_state = { mutex: Mutex.new, stacks: {} }
+            lifted = lambda do |*flat_vars|
+              state = grad_call_state_current(call_state)
+              if state.nil?
+                raise RuntimeError, "gradient transform invoked without call state"
+              end
+
+              call_args, call_kwargs = apply_flat_vars_to_targets(
+                state[:args],
+                state[:kwargs],
+                state[:selections],
+                flat_vars
+              )
+              raw_value = fun.call(*call_args, **call_kwargs)
+              state[:captured_value] = raw_value
+              extract_loss(raw_value)
+            end
+
+            native_argnums = (0...flat_inputs.length).to_a
+            native_fn = if with_value
+              native_value_and_grad(lifted, native_argnums)
+            else
+              native_grad(lifted, native_argnums)
+            end
+
+            entry = {
+              native_fn: native_fn,
+              call_state: call_state
+            }
+            cache[cache_key] = entry
           end
 
+          state = {
+            args: args,
+            kwargs: kwargs,
+            selections: selections,
+            captured_value: nil
+          }
+          grad_call_state_push(entry[:call_state], state)
+
           if with_value
-            native_fn = native_value_and_grad(lifted, native_argnums)
-            _loss, raw_grads = native_fn.call(*flat_inputs)
-            value = captured_value.nil? ? fun.call(*args, **kwargs) : captured_value
+            _loss, raw_grads = entry[:native_fn].call(*flat_inputs)
+            value = state[:captured_value]
+            value = fun.call(*args, **kwargs) if value.nil?
             [value, rebuild_grad_result(raw_grads, selections, argnames)]
           else
-            native_fn = native_grad(lifted, native_argnums)
-            raw_grads = native_fn.call(*flat_inputs)
+            raw_grads = entry[:native_fn].call(*flat_inputs)
             rebuild_grad_result(raw_grads, selections, argnames)
           end
+        ensure
+          grad_call_state_pop(entry[:call_state]) unless entry.nil?
+        end
+      end
+
+      def grad_call_state_current(call_state)
+        thread = Thread.current
+        call_state[:mutex].synchronize do
+          stack = call_state[:stacks][thread]
+          stack&.last
+        end
+      end
+
+      def grad_call_state_push(call_state, state)
+        thread = Thread.current
+        call_state[:mutex].synchronize do
+          stack = call_state[:stacks][thread]
+          if stack.nil?
+            stack = []
+            call_state[:stacks][thread] = stack
+          end
+          stack << state
+        end
+      end
+
+      def grad_call_state_pop(call_state)
+        thread = Thread.current
+        call_state[:mutex].synchronize do
+          stack = call_state[:stacks][thread]
+          return if stack.nil?
+
+          stack.pop
+          call_state[:stacks].delete(thread) if stack.empty?
         end
       end
 
@@ -725,6 +909,119 @@ module MLX
         end
       end
 
+      def extract_loss(output)
+        return output if output.is_a?(MLX::Core::Array)
+
+        if output.is_a?(::Array) && !output.empty? && output[0].is_a?(MLX::Core::Array)
+          return output[0]
+        end
+
+        raise ArgumentError,
+              "function must return an MLX::Core::Array or an Array whose first element is an MLX::Core::Array"
+      end
+
+      def build_target_selections(args, kwargs, argnums, argnames)
+        positional = []
+        keyword = []
+        flat_inputs = []
+
+        argnums.each do |index|
+          if index >= args.length
+            raise ArgumentError,
+                  "Can't compute gradient for positional argument #{index} when #{args.length} positional arguments were provided"
+          end
+          spec = flatten_tree_spec(args[index], flat_inputs, true)
+          positional << { index: index, spec: spec }
+        end
+
+        argnames.each do |name|
+          key = kwarg_key_for_name(kwargs, name)
+          unless key
+            raise ArgumentError,
+                  "Can't compute gradient for keyword argument '#{name}' because it was not provided"
+          end
+          spec = flatten_tree_spec(kwargs[key], flat_inputs, true)
+          keyword << { key: key, name: name, spec: spec }
+        end
+
+        [{ positional: positional, keyword: keyword }, flat_inputs]
+      end
+
+      def grad_selection_cache_key(selections)
+        positional = selections[:positional].map do |entry|
+          "#{entry[:index]}:#{structure_cache_key(entry[:spec])}"
+        end
+        keyword = selections[:keyword].map do |entry|
+          "#{entry[:name]}:#{entry[:key]}:#{structure_cache_key(entry[:spec])}"
+        end
+        "P[#{positional.join(',')}]K[#{keyword.join(',')}]"
+      end
+
+      def normalize_raw_grads(raw)
+        normalize_array_sequence(raw, "gradient")
+      end
+
+      def rebuild_grad_result(raw_grads, selections, argnames)
+        grad_arrays = normalize_raw_grads(raw_grads)
+        cursor = 0
+
+        positional_grads = selections[:positional].map do |entry|
+          value, cursor = inflate_tree_from_arrays(entry[:spec], grad_arrays, cursor)
+          value
+        end
+        keyword_grads = {}
+        selections[:keyword].each do |entry|
+          value, cursor = inflate_tree_from_arrays(entry[:spec], grad_arrays, cursor)
+          keyword_grads[entry[:name]] = value
+        end
+        unless cursor == grad_arrays.length
+          raise RuntimeError, "internal gradient reconstruction mismatch"
+        end
+
+        if argnames.empty?
+          return positional_grads[0] if positional_grads.length == 1
+          return positional_grads
+        end
+
+        positional_out = if positional_grads.empty?
+          nil
+        elsif positional_grads.length == 1
+          positional_grads[0]
+        else
+          positional_grads
+        end
+        [positional_out, keyword_grads]
+      end
+
+      def apply_flat_vars_to_targets(args, kwargs, selections, flat_vars)
+        rebuilt_args = args.dup
+        rebuilt_kwargs = kwargs.dup
+        cursor = 0
+
+        selections[:positional].each do |entry|
+          value, cursor = inflate_tree_from_arrays(entry[:spec], flat_vars, cursor)
+          rebuilt_args[entry[:index]] = value
+        end
+
+        selections[:keyword].each do |entry|
+          value, cursor = inflate_tree_from_arrays(entry[:spec], flat_vars, cursor)
+          rebuilt_kwargs[entry[:key]] = value
+        end
+
+        unless cursor == flat_vars.length
+          raise RuntimeError, "internal target reconstruction mismatch"
+        end
+        [rebuilt_args, rebuilt_kwargs]
+      end
+
+      def kwarg_key_for_name(kwargs, name)
+        symbol = name.to_sym
+        return symbol if kwargs.key?(symbol)
+        return name if kwargs.key?(name)
+
+        nil
+      end
+
       def custom_jvp(fun, primals, tangents)
         primals_list = normalize_array_output(primals, "primals")
         tangents_list = normalize_array_output(tangents, "tangents")
@@ -768,44 +1065,6 @@ module MLX
         end
       end
 
-      def extract_loss(output)
-        return output if output.is_a?(MLX::Core::Array)
-
-        if output.is_a?(::Array) && !output.empty? && output[0].is_a?(MLX::Core::Array)
-          return output[0]
-        end
-
-        raise ArgumentError,
-              "function must return an MLX::Core::Array or an Array whose first element is an MLX::Core::Array"
-      end
-
-      def build_target_selections(args, kwargs, argnums, argnames)
-        positional = []
-        keyword = []
-        flat_inputs = []
-
-        argnums.each do |index|
-          if index >= args.length
-            raise ArgumentError,
-                  "Can't compute gradient for positional argument #{index} when #{args.length} positional arguments were provided"
-          end
-          spec = flatten_tree_spec(args[index], flat_inputs, true)
-          positional << { index: index, spec: spec }
-        end
-
-        argnames.each do |name|
-          key = kwarg_key_for_name(kwargs, name)
-          unless key
-            raise ArgumentError,
-                  "Can't compute gradient for keyword argument '#{name}' because it was not provided"
-          end
-          spec = flatten_tree_spec(kwargs[key], flat_inputs, true)
-          keyword << { key: key, name: name, spec: spec }
-        end
-
-        [{ positional: positional, keyword: keyword }, flat_inputs]
-      end
-
       def flatten_tree_spec(value, arrays, strict_arrays)
         if value.is_a?(MLX::Core::Array)
           arrays << value
@@ -873,10 +1132,6 @@ module MLX
         end
       end
 
-      def normalize_raw_grads(raw)
-        normalize_array_sequence(raw, "gradient")
-      end
-
       def normalize_array_sequence(raw, context)
         return [raw] if raw.is_a?(MLX::Core::Array)
 
@@ -896,66 +1151,6 @@ module MLX
         end
       end
 
-      def rebuild_grad_result(raw_grads, selections, argnames)
-        grad_arrays = normalize_raw_grads(raw_grads)
-        cursor = 0
-
-        positional_grads = selections[:positional].map do |entry|
-          value, cursor = inflate_tree_from_arrays(entry[:spec], grad_arrays, cursor)
-          value
-        end
-        keyword_grads = {}
-        selections[:keyword].each do |entry|
-          value, cursor = inflate_tree_from_arrays(entry[:spec], grad_arrays, cursor)
-          keyword_grads[entry[:name]] = value
-        end
-        unless cursor == grad_arrays.length
-          raise RuntimeError, "internal gradient reconstruction mismatch"
-        end
-
-        if argnames.empty?
-          return positional_grads[0] if positional_grads.length == 1
-          return positional_grads
-        end
-
-        positional_out = if positional_grads.empty?
-          nil
-        elsif positional_grads.length == 1
-          positional_grads[0]
-        else
-          positional_grads
-        end
-        [positional_out, keyword_grads]
-      end
-
-      def apply_flat_vars_to_targets(args, kwargs, selections, flat_vars)
-        rebuilt_args = args.dup
-        rebuilt_kwargs = kwargs.dup
-        cursor = 0
-
-        selections[:positional].each do |entry|
-          value, cursor = inflate_tree_from_arrays(entry[:spec], flat_vars, cursor)
-          rebuilt_args[entry[:index]] = value
-        end
-
-        selections[:keyword].each do |entry|
-          value, cursor = inflate_tree_from_arrays(entry[:spec], flat_vars, cursor)
-          rebuilt_kwargs[entry[:key]] = value
-        end
-
-        unless cursor == flat_vars.length
-          raise RuntimeError, "internal target reconstruction mismatch"
-        end
-        [rebuilt_args, rebuilt_kwargs]
-      end
-
-      def kwarg_key_for_name(kwargs, name)
-        symbol = name.to_sym
-        return symbol if kwargs.key?(symbol)
-        return name if kwargs.key?(name)
-
-        nil
-      end
     end
 
     class Device
@@ -1034,8 +1229,8 @@ module MLX
         MLX::Core.cos(self)
       end
 
-      def mean(axis = nil)
-        MLX::Core.mean(self, axis)
+      def mean(axis = nil, keepdims_positional = nil, keepdims: nil)
+        MLX::Core.mean(self, axis, keepdims_positional, keepdims: keepdims)
       end
 
       def sum(axis = nil)
@@ -1307,6 +1502,10 @@ module MLX
         MLX::Core.negative(self)
       end
 
+      def -@
+        __neg__
+      end
+
       def __pow__(other)
         MLX::Core.power(self, other)
       end
@@ -1375,18 +1574,34 @@ module MLX
         MLX::Core.less(self, other)
       end
 
+      def <(other)
+        __lt__(other)
+      end
+
       def __le__(other)
         MLX::Core.less_equal(self, other)
       end
 
+      def <=(other)
+        __le__(other)
+      end
+
       def __gt__(other)
         MLX::Core.greater(self, other)
       end
 
+      def >(other)
+        __gt__(other)
+      end
+
       def __ge__(other)
         MLX::Core.greater_equal(self, other)
       end
 
+      def >=(other)
+        __ge__(other)
+      end
+
       def __iadd__(other)
         __add__(other)
       end
@@ -1439,6 +1654,16 @@ module MLX
         MLX::Core.floor_divide(other, self)
       end
 
+      def coerce(other)
+        if other.is_a?(MLX::Core::Array)
+          [other, self]
+        elsif other.is_a?(::Numeric)
+          [MLX::Core.array(other, dtype), self]
+        else
+          raise TypeError, "#{other.class} can't be coerced into MLX::Core::Array"
+        end
+      end
+
       def __getitem__(index)
         self[index]
       end

data/lib/mlx/distributed_utils/common.rb

@@ -5,7 +5,7 @@ require "json"
 
 module MLX
   module DistributedUtils
-    Host = Struct.new(:rank, :ssh_hostname, :ips, :rdma, keyword_init: true)
+    Host = Data.define(:rank, :ssh_hostname, :ips, :rdma)
 
     class Hostfile
       attr_accessor :hosts, :backend, :envs

data/lib/mlx/distributed_utils/config.rb

@@ -8,13 +8,14 @@ require "shellwords"
 
 module MLX
   module DistributedUtils
-    SSHInfo = Struct.new(:can_ssh, :has_sudo, keyword_init: true) do
+    SSHInfo = Data.define(:can_ssh, :has_sudo) do
       def to_bool
         can_ssh
       end
     end
-    ThunderboltPort = Struct.new(:iface, :uuid, :connected_to, keyword_init: true)
-    ThunderboltHost = Struct.new(:name, :ports, keyword_init: true)
+    ThunderboltPort = Data.define(:iface, :uuid, :connected_to)
+    ThunderboltHost = Data.define(:name, :ports)
+    CommandResult = Data.define(:stdout, :stderr, :status)
 
     class IPConfigurator
       attr_reader :ips, :hosts, :tb_hosts
@@ -509,6 +510,8 @@ module MLX
       end
 
       def config_main(argv = ARGV, runner: nil)
+        Process.warmup if Process.respond_to?(:warmup)
+
         opts = {
           verbose: false,
           hosts: "127.0.0.1",
@@ -577,7 +580,7 @@ module MLX
         return runner.call(cmd) unless runner.nil?
 
         stdout, stderr, status = Open3.capture3(*cmd)
-        Struct.new(:stdout, :stderr, :status, keyword_init: true).new(stdout: stdout, stderr: stderr, status: status)
+        CommandResult.new(stdout: stdout, stderr: stderr, status: status)
       end
 
       def stdout_for(result)

data/lib/mlx/distributed_utils/launch.rb

@@ -314,6 +314,8 @@ module MLX
       end
 
       def main(argv = ARGV)
+        Process.warmup if Process.respond_to?(:warmup)
+
         opts = {
           print_python: false,
           verbose: false,