mlx 0.30.7.2 → 0.30.7.6

data/lib/mlx/dsl/tensor.rb

@@ -0,0 +1,126 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    module Tensor
+      module_function
+
+      def scatter_rows(base:, row_indices:, values:, axis: nil)
+        axis = __dsl_default_scatter_axis(base, axis)
+        row_indices = __dsl_to_index_array(row_indices)
+        values = __dsl_to_array(values, dtype: base.dtype)
+
+        if base.ndim != values.ndim
+          raise ArgumentError, "base and values must have the same rank"
+        end
+
+        unless base.shape.each_with_index.all? { |dim, idx| idx == axis || dim == values.shape[idx] }
+          raise ArgumentError, "values shape must match base shape except along axis #{axis}"
+        end
+
+        indices = __dsl_expand_indices_for_axis(
+          row_indices: row_indices,
+          values_shape: values.shape,
+          axis: axis
+        )
+        MLX::Core.put_along_axis(base, indices, values, axis)
+      end
+
+      def where_labels(base:, labels:, mapping:, mode: :add_or_replace)
+        mode = mode.to_sym
+        unless [:add_or_replace, :replace].include?(mode)
+          raise ArgumentError, "mode must be :add_or_replace or :replace"
+        end
+
+        out = base
+        mapping.each do |label_value, mapped_value|
+          mask = MLX::Core.equal(labels, label_value)
+          mask = __dsl_expand_trailing_dims(mask, out.ndim)
+          mapped = __dsl_broadcast_mapping(mapped_value, base: out, labels_ndim: labels.ndim)
+          replacement = if mode == :replace
+            mapped
+          else
+            MLX::Core.add(out, mapped)
+          end
+          out = MLX::Core.where(mask, replacement, out)
+        end
+        out
+      end
+
+      def __dsl_default_scatter_axis(base, axis)
+        return axis.to_i unless axis.nil?
+        return 1 if base.ndim == 3
+
+        0
+      end
+      private_class_method :__dsl_default_scatter_axis
+
+      def __dsl_to_index_array(indices)
+        if indices.is_a?(MLX::Core::Array)
+          return indices.astype(MLX::Core.int32)
+        end
+
+        MLX::Core.array(indices, MLX::Core.int32)
+      end
+      private_class_method :__dsl_to_index_array
+
+      def __dsl_to_array(value, dtype:)
+        if value.is_a?(MLX::Core::Array)
+          return value.astype(dtype)
+        end
+
+        MLX::Core.array(value, dtype)
+      end
+      private_class_method :__dsl_to_array
+
+      def __dsl_expand_indices_for_axis(row_indices:, values_shape:, axis:)
+        if row_indices.ndim == 1
+          expected = values_shape[axis]
+          if row_indices.shape[0] != expected
+            raise ArgumentError, "row_indices length must match values shape at axis #{axis}"
+          end
+
+          reshape = Array.new(values_shape.length, 1)
+          reshape[axis] = expected
+          base = MLX::Core.reshape(row_indices, reshape)
+          return MLX::Core.broadcast_to(base, values_shape)
+        end
+
+        if row_indices.shape == values_shape
+          return row_indices
+        end
+
+        raise ArgumentError, "row_indices must be rank-1 or match values shape"
+      end
+      private_class_method :__dsl_expand_indices_for_axis
+
+      def __dsl_expand_trailing_dims(array, target_ndim)
+        out = array
+        while out.ndim < target_ndim
+          out = MLX::Core.expand_dims(out, out.ndim)
+        end
+        out
+      end
+      private_class_method :__dsl_expand_trailing_dims
+
+      def __dsl_broadcast_mapping(value, base:, labels_ndim:)
+        mapped = __dsl_to_array(value, dtype: base.dtype)
+        return mapped if mapped.shape == base.shape
+
+        if mapped.ndim == 1 && mapped.shape[0] == base.shape[-1]
+          reshape = Array.new(labels_ndim, 1) + [mapped.shape[0]]
+          mapped = MLX::Core.reshape(mapped, reshape)
+          return MLX::Core.broadcast_to(mapped, base.shape)
+        end
+
+        if mapped.ndim == labels_ndim && mapped.shape == base.shape[0...labels_ndim]
+          mapped = __dsl_expand_trailing_dims(mapped, base.ndim)
+          return MLX::Core.broadcast_to(mapped, base.shape)
+        end
+
+        MLX::Core.broadcast_to(mapped, base.shape)
+      end
+      private_class_method :__dsl_broadcast_mapping
+    end
+  end
+end

data/lib/mlx/dsl/transformer_block.rb

@@ -0,0 +1,113 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class FeedForward < MLX::NN::Module
+      def initialize(dims:, hidden_dims:, kind: :gelu, bias: false)
+        super()
+        @kind = kind.to_sym
+        @dims = Integer(dims)
+        @hidden_dims = Integer(hidden_dims)
+
+        case @kind
+        when :swiglu
+          self.gate_proj = MLX::NN::Linear.new(@dims, @hidden_dims, bias: bias)
+          self.up_proj = MLX::NN::Linear.new(@dims, @hidden_dims, bias: bias)
+          self.down_proj = MLX::NN::Linear.new(@hidden_dims, @dims, bias: bias)
+        else
+          self.in_proj = MLX::NN::Linear.new(@dims, @hidden_dims, bias: bias)
+          self.out_proj = MLX::NN::Linear.new(@hidden_dims, @dims, bias: bias)
+        end
+      end
+
+      def call(x)
+        case @kind
+        when :swiglu
+          gated = MLX::NN.silu(gate_proj.call(x))
+          down_proj.call(MLX::Core.multiply(gated, up_proj.call(x)))
+        when :relu
+          out_proj.call(MLX::NN.relu(in_proj.call(x)))
+        else
+          out_proj.call(MLX::NN.gelu(in_proj.call(x)))
+        end
+      end
+    end
+
+    class TransformerBlock < MLX::NN::Module
+      def initialize(
+        dims:,
+        num_heads:,
+        kv_heads: nil,
+        ffn_dims: nil,
+        norm: :rms,
+        norm_eps: 1e-5,
+        ffn: nil,
+        rope: nil,
+        qkv_bias: false,
+        backend: :sdpa,
+        cache: false
+      )
+        super()
+
+        ffn_config = (ffn || {}).transform_keys(&:to_sym)
+        ffn_kind = ffn_config.fetch(:kind, :gelu)
+        hidden_dims = ffn_dims || ffn_config.fetch(:hidden_dims, Integer(dims) * 4)
+        ffn_bias = ffn_config.fetch(:bias, false)
+
+        self.attention_norm = __dsl_build_norm(norm, Integer(dims), norm_eps)
+        self.attention = MLX::DSL::Attention.new(
+          dims: dims,
+          num_heads: num_heads,
+          kv_heads: kv_heads,
+          qkv_bias: qkv_bias,
+          backend: backend,
+          rope: rope,
+          cache: cache
+        )
+        self.ffn_norm = __dsl_build_norm(norm, Integer(dims), norm_eps)
+        self.feed_forward = MLX::DSL::FeedForward.new(
+          dims: dims,
+          hidden_dims: hidden_dims,
+          kind: ffn_kind,
+          bias: ffn_bias
+        )
+
+        @cache_enabled = !!cache
+      end
+
+      def call(x, mask: nil, cache: nil, **_kwargs)
+        attn_input = attention_norm.call(x)
+        attn_result = attention.call(attn_input, attn_input, attn_input, mask: mask, cache: cache)
+        if attn_result.is_a?(Array) && attn_result.length == 2
+          attn_out, next_cache = attn_result
+        else
+          attn_out = attn_result
+          next_cache = cache
+        end
+
+        hidden = MLX::Core.add(x, attn_out)
+        ffn_out = feed_forward.call(ffn_norm.call(hidden))
+        output = MLX::Core.add(hidden, ffn_out)
+
+        if @cache_enabled || !cache.nil?
+          [output, next_cache]
+        else
+          output
+        end
+      end
+
+      private
+
+      def __dsl_build_norm(kind, dims, eps)
+        case kind.to_sym
+        when :layer, :layer_norm
+          MLX::NN::LayerNorm.new(dims, eps: eps)
+        when :rms, :rms_norm
+          MLX::NN::RMSNorm.new(dims, eps: eps)
+        else
+          raise ArgumentError, "unsupported norm kind: #{kind.inspect}"
+        end
+      end
+    end
+  end
+end

data/lib/mlx/dsl/weight_map.rb

@@ -0,0 +1,140 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    def self.weight_map(&block)
+      mapper = WeightMap.new
+      mapper.instance_eval(&block) if block_given?
+      mapper
+    end
+
+    class WeightMap
+      def initialize
+        @rules = []
+      end
+
+      def strip_prefix(prefix)
+        @rules << [:strip_prefix, prefix.to_s]
+        self
+      end
+
+      def rename(from = nil, to = nil)
+        if from.is_a?(Hash)
+          from.each do |src, dst|
+            @rules << [:rename, src.to_s, dst.to_s]
+          end
+          return self
+        end
+
+        if from.nil? || to.nil?
+          raise ArgumentError, "rename requires either a mapping hash or from/to arguments"
+        end
+
+        @rules << [:rename, from.to_s, to.to_s]
+        self
+      end
+
+      def regex(pattern, replacement = nil, &block)
+        if pattern.nil?
+          raise ArgumentError, "regex requires a Regexp pattern"
+        end
+        if replacement.nil? && !block_given?
+          raise ArgumentError, "regex requires a replacement argument or block"
+        end
+
+        @rules << [:regex, pattern, replacement, block]
+        self
+      end
+
+      def split_qkv(source, into:, axis: 0)
+        names = into.to_a.map(&:to_s)
+        if names.empty? || names.length < 2
+          raise ArgumentError, "split_qkv :into must include at least two output names"
+        end
+
+        @rules << [:split, source.to_s, names, axis.to_i]
+        self
+      end
+
+      def transpose_if(rank:, order:)
+        @rules << [:transpose_if, rank.to_i, order.to_a]
+        self
+      end
+
+      def apply(weights)
+        entries = if weights.is_a?(Hash)
+          weights.to_a
+        elsif weights.respond_to?(:to_a)
+          weights.to_a
+        else
+          raise ArgumentError, "weights must be a Hash or array-like key/value collection"
+        end
+
+        out = {}
+        entries.each do |entry|
+          key, value = entry
+          __dsl_apply_rules_for_entry(key.to_s, value).each do |mapped_key, mapped_value|
+            out[mapped_key] = mapped_value
+          end
+        end
+        out
+      end
+
+      private
+
+      def __dsl_apply_rules_for_entry(key, value)
+        current_entries = [[key, value]]
+        @rules.each do |rule|
+          current_entries = current_entries.flat_map do |curr_key, curr_value|
+            __dsl_apply_rule(rule, curr_key, curr_value)
+          end
+        end
+        current_entries
+      end
+
+      def __dsl_apply_rule(rule, key, value)
+        kind = rule[0]
+        case kind
+        when :strip_prefix
+          prefix = rule[1]
+          if key.start_with?(prefix)
+            [[key[prefix.length..], value]]
+          else
+            [[key, value]]
+          end
+        when :rename
+          from = rule[1]
+          to = rule[2]
+          [[key.gsub(from, to), value]]
+        when :regex
+          pattern = rule[1]
+          replacement = rule[2]
+          block = rule[3]
+          if block.nil?
+            [[key.gsub(pattern, replacement.to_s), value]]
+          else
+            [[key.gsub(pattern, &block), value]]
+          end
+        when :split
+          source = rule[1]
+          targets = rule[2]
+          axis = rule[3]
+          return [[key, value]] unless key == source
+
+          parts = MLX::Core.split(value, targets.length, axis)
+          targets.each_with_index.map { |target, index| [target, parts[index]] }
+        when :transpose_if
+          target_rank = rule[1]
+          order = rule[2]
+          if value.respond_to?(:shape) && value.shape.length == target_rank
+            [[key, MLX::Core.transpose(value, order)]]
+          else
+            [[key, value]]
+          end
+        else
+          [[key, value]]
+        end
+      end
+    end
+  end
+end

data/lib/mlx/dsl.rb

@@ -10,6 +10,16 @@ require_relative "dsl/data_pipeline"
 require_relative "dsl/experiment"
 require_relative "dsl/split_plan"
 require_relative "dsl/builder"
+require_relative "dsl/config_schema"
+require_relative "dsl/weight_map"
+require_relative "dsl/kv_cache"
+require_relative "dsl/masks"
+require_relative "dsl/positions"
+require_relative "dsl/tensor"
+require_relative "dsl/run_stack"
+require_relative "dsl/attention"
+require_relative "dsl/transformer_block"
+require_relative "dsl/generate"
 require_relative "dsl/train_step"
 require_relative "dsl/model_mixin"
 require_relative "dsl/model"

data/lib/mlx/nn/base.rb

@@ -323,6 +323,8 @@ module MLX
               current_value = dst[k]
               if current_value.is_a?(Module) && new_value.is_a?(Module)
                 dst[k] = new_value
+              elsif current_value.is_a?(Module) && (new_value.is_a?(Hash) || new_value.is_a?(Array))
+                update_modules_impl(current_value, new_value, strict)
               elsif current_value.is_a?(Hash) || current_value.is_a?(Array)
                 update_modules_impl(current_value, new_value, strict)
               elsif strict && new_value != {}
@@ -337,6 +339,8 @@ module MLX
             current_value = dst[i]
             if current_value.is_a?(Module) && new_value.is_a?(Module)
               dst[i] = new_value
+            elsif current_value.is_a?(Module) && (new_value.is_a?(Hash) || new_value.is_a?(Array))
+              update_modules_impl(current_value, new_value, strict)
             elsif current_value.is_a?(Hash) || current_value.is_a?(Array)
               update_modules_impl(current_value, new_value, strict)
             elsif strict && new_value != {}

data/lib/mlx/nn/layers/linear.rb

@@ -23,11 +23,10 @@ module MLX
       end
 
       def call(x)
-        out = MLX::Core.matmul(x, weight.T)
         if state.key?("bias")
-          MLX::Core.add(out, bias)
+          MLX::Core.addmm(bias, x, weight.T)
         else
-          out
+          MLX::Core.matmul(x, weight.T)
         end
       end
 

data/lib/mlx/onnx.rb

@@ -0,0 +1,250 @@
+# frozen_string_literal: true
+
+require "json"
+require_relative "../mlx-onnx/webgpu_harness"
+
+module MLX
+  module ONNX
+    ONNX_VERSION = 1
+
+    module_function
+
+    def export_onnx(
+      target_path,
+      fun,
+      *extras,
+      shapeless: false,
+      opset: 18,
+      model_name: "mlx_graph",
+      external_data: false,
+      external_data_file: nil,
+      external_data_size_threshold: 1024,
+      **trace_kwargs
+    )
+      ensure_native_graph_ir!
+      assert_boolean!(shapeless, "shapeless")
+      assert_boolean!(external_data, "external_data")
+      native_target_path = normalize_binary_target_path!(target_path, "export_onnx")
+      native_fun, native_extras, _positional_count, _keyword_names =
+        prepare_trace_invocation_without_native_kwargs(fun, extras, trace_kwargs)
+      translate_native_unsupported do
+        written = MLX::ONNX::Native.export_onnx(
+          native_target_path,
+          native_fun,
+          native_extras,
+          {},
+          shapeless,
+          opset,
+          model_name,
+          external_data,
+          external_data_file,
+          external_data_size_threshold
+        )
+        unless written.is_a?(String)
+          raise TypeError, "MLX::ONNX::Native.export_onnx must return String path"
+        end
+        written
+      end
+    end
+
+    def export_onnx_json(
+      fun,
+      *extras,
+      shapeless: false,
+      opset: 18,
+      model_name: "mlx_graph",
+      **trace_kwargs
+    )
+      ensure_native_graph_ir!
+      assert_boolean!(shapeless, "shapeless")
+      native_fun, native_extras, _positional_count, _keyword_names =
+        prepare_trace_invocation_without_native_kwargs(fun, extras, trace_kwargs)
+      translate_native_unsupported do
+        onnx_json = MLX::ONNX::Native.export_onnx_json(
+          native_fun,
+          native_extras,
+          {},
+          shapeless,
+          opset,
+          model_name
+        )
+        unless onnx_json.is_a?(String)
+          raise TypeError, "MLX::ONNX::Native.export_onnx_json must return String JSON"
+        end
+        onnx_json
+      end
+    end
+
+    def export_onnx_compatibility_report(fun, *extras, shapeless: false, **trace_kwargs)
+      ensure_native_graph_ir!
+      assert_boolean!(shapeless, "shapeless")
+      native_fun, native_extras, _positional_count, _keyword_names =
+        prepare_trace_invocation_without_native_kwargs(fun, extras, trace_kwargs)
+      report = MLX::ONNX::Native.export_onnx_compatibility_report(
+        native_fun,
+        native_extras,
+        {},
+        shapeless
+      )
+      unless report.is_a?(Hash)
+        raise TypeError, "MLX::ONNX::Native.export_onnx_compatibility_report must return Hash payload"
+      end
+
+      report
+    end
+
+    def export_graph_ir(fun, *extras, shapeless: false, **trace_kwargs)
+      ensure_native_graph_ir!
+      assert_boolean!(shapeless, "shapeless")
+      native_fun, native_extras, positional_count, keyword_names =
+        prepare_trace_invocation_without_native_kwargs(fun, extras, trace_kwargs)
+      payload = MLX::ONNX::Native.export_graph_ir(native_fun, native_extras, {}, shapeless)
+      unless payload.is_a?(Hash)
+        raise TypeError, "MLX::ONNX::Native.export_graph_ir must return Hash payload"
+      end
+      inject_keyword_inputs!(payload, positional_count, keyword_names)
+      payload
+    end
+
+    def export_graph_ir_json(fun, *extras, shapeless: false, **trace_kwargs)
+      ensure_native_graph_ir!
+      assert_boolean!(shapeless, "shapeless")
+      native_fun, native_extras, positional_count, keyword_names =
+        prepare_trace_invocation_without_native_kwargs(fun, extras, trace_kwargs)
+      content = MLX::ONNX::Native.export_graph_ir_json(native_fun, native_extras, {}, shapeless)
+      unless content.is_a?(String)
+        raise TypeError, "MLX::ONNX::Native.export_graph_ir_json must return String JSON"
+      end
+      return content if keyword_names.empty?
+
+      payload = JSON.parse(content)
+      inject_keyword_inputs!(payload, positional_count, keyword_names)
+      content = JSON.generate(payload)
+      content
+    end
+
+    def graph_ir_to_onnx(
+      target_path,
+      ir_source,
+      opset: 18,
+      model_name: "mlx_graph",
+      external_data: false,
+      external_data_file: nil,
+      external_data_size_threshold: 1024
+    )
+      ensure_native_graph_ir!
+      assert_boolean!(external_data, "external_data")
+      native_target_path = normalize_binary_target_path!(target_path, "graph_ir_to_onnx")
+      translate_native_unsupported do
+        written = MLX::ONNX::Native.graph_ir_to_onnx(
+          native_target_path,
+          ir_source,
+          opset,
+          model_name,
+          external_data,
+          external_data_file,
+          external_data_size_threshold
+        )
+        unless written.is_a?(String)
+          raise TypeError, "MLX::ONNX::Native.graph_ir_to_onnx must return String path"
+        end
+        written
+      end
+    end
+
+    def graph_ir_to_onnx_json(ir_source, opset: 18, model_name: "mlx_graph")
+      ensure_native_graph_ir!
+      translate_native_unsupported do
+        onnx_json = MLX::ONNX::Native.graph_ir_to_onnx_json(ir_source, opset, model_name)
+        unless onnx_json.is_a?(String)
+          raise TypeError, "MLX::ONNX::Native.graph_ir_to_onnx_json must return String JSON"
+        end
+        onnx_json
+      end
+    end
+
+    def ensure_native_graph_ir!
+      MLX::Core.ensure_native!
+      unless defined?(MLX::ONNX::Native)
+        raise RuntimeError, "MLX::ONNX::Native is unavailable"
+      end
+    end
+    private_class_method :ensure_native_graph_ir!
+
+    def assert_boolean!(value, label)
+      unless value == true || value == false
+        raise TypeError, "#{label} must be true or false"
+      end
+    end
+    private_class_method :assert_boolean!
+
+    def translate_native_unsupported
+      yield
+    rescue MLX::ONNX::Native::UnsupportedError => e
+      raise NotImplementedError, e.message
+    end
+    private_class_method :translate_native_unsupported
+
+    def normalize_binary_target_path!(target_path, method_name)
+      if target_path.respond_to?(:write) && !target_path.respond_to?(:to_path)
+        raise ArgumentError, "#{method_name} requires a path-like target, not an IO-like target"
+      end
+
+      path = if target_path.respond_to?(:to_path)
+        target_path.to_path
+      else
+        target_path
+      end
+
+      unless path.is_a?(String)
+        raise TypeError, "target_path must be a String or respond to #to_path"
+      end
+
+      path
+    end
+    private_class_method :normalize_binary_target_path!
+
+    def prepare_trace_invocation_without_native_kwargs(fun, extras, trace_kwargs)
+      positional_inputs = extras.dup
+      return [fun, positional_inputs, positional_inputs.length, []] if trace_kwargs.empty?
+
+      keyword_entries = trace_kwargs.to_a
+      positional_count = positional_inputs.length
+      native_fun = lambda do |*all_args|
+        positional_args = all_args.first(positional_count)
+        keyword_args = {}
+        keyword_entries.each_with_index do |(name, _value), idx|
+          keyword_args[name] = all_args[positional_count + idx]
+        end
+        fun.call(*positional_args, **keyword_args)
+      end
+
+      keyword_values = keyword_entries.map { |(_name, value)| value }
+      keyword_names = keyword_entries.map { |(name, _value)| name.to_s }
+      [native_fun, positional_inputs + keyword_values, positional_count, keyword_names]
+    end
+    private_class_method :prepare_trace_invocation_without_native_kwargs
+
+    def inject_keyword_inputs!(payload, positional_count, keyword_names)
+      return payload if keyword_names.empty?
+      return payload unless payload.is_a?(Hash)
+
+      inputs = payload["inputs"]
+      return payload unless inputs.is_a?(Array)
+
+      keyword_inputs = []
+      keyword_names.each_with_index do |name, idx|
+        input_entry = inputs[positional_count + idx]
+        next unless input_entry.is_a?(Hash)
+
+        tensor = input_entry["name"]
+        next unless tensor.is_a?(String) && !tensor.empty?
+
+        keyword_inputs << { "name" => name, "tensor" => tensor }
+      end
+      payload["keyword_inputs"] = keyword_inputs
+      payload
+    end
+    private_class_method :inject_keyword_inputs!
+  end
+end

data/lib/mlx/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module MLX
-  VERSION = "0.30.7.2"
+  VERSION = "0.30.7.6"
 end