mlx 0.30.7.2 → 0.30.7.6

data/lib/mlx/dsl/attention.rb

@@ -0,0 +1,132 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class Attention < MLX::NN::Module
+      def initialize(
+        dims:,
+        num_heads:,
+        kv_heads: nil,
+        qkv_bias: false,
+        backend: :sdpa,
+        rope: nil,
+        cache: false
+      )
+        super()
+
+        @dims = Integer(dims)
+        @num_heads = Integer(num_heads)
+        @kv_heads = kv_heads.nil? ? @num_heads : Integer(kv_heads)
+        if (@dims % @num_heads) != 0
+          raise ArgumentError, "dims must be divisible by num_heads"
+        end
+        if (@num_heads % @kv_heads) != 0
+          raise ArgumentError, "num_heads must be divisible by kv_heads"
+        end
+
+        @head_dim = @dims / @num_heads
+        @kv_repeats = @num_heads / @kv_heads
+        @backend = backend.to_sym
+        @cache_enabled = !!cache
+        @scale = Math.sqrt(1.0 / @head_dim)
+
+        self.query_proj = MLX::NN::Linear.new(@dims, @num_heads * @head_dim, bias: qkv_bias)
+        self.key_proj = MLX::NN::Linear.new(@dims, @kv_heads * @head_dim, bias: qkv_bias)
+        self.value_proj = MLX::NN::Linear.new(@dims, @kv_heads * @head_dim, bias: qkv_bias)
+        self.out_proj = MLX::NN::Linear.new(@num_heads * @head_dim, @dims, bias: qkv_bias)
+        self.rope = __dsl_build_rope(rope)
+      end
+
+      def call(queries, keys = nil, values = nil, mask: nil, cache: nil)
+        keys ||= queries
+        values ||= keys
+        q_was_2d = queries.ndim == 2
+
+        queries = MLX::Core.expand_dims(queries, 0) if q_was_2d
+        keys = MLX::Core.expand_dims(keys, 0) if keys.ndim == 2
+        values = MLX::Core.expand_dims(values, 0) if values.ndim == 2
+
+        batch_size, q_len, = queries.shape
+
+        q = __dsl_pack_heads(query_proj.call(queries), @num_heads)
+        k = __dsl_pack_heads(key_proj.call(keys), @kv_heads)
+        v = __dsl_pack_heads(value_proj.call(values), @kv_heads)
+
+        offset = cache.nil? ? 0 : cache[0].shape[2]
+        if !rope.nil?
+          if offset.zero?
+            q = rope.call(q)
+            k = rope.call(k)
+          else
+            q = rope.call(q, offset: offset)
+            k = rope.call(k, offset: offset)
+          end
+        end
+
+        unless cache.nil?
+          key_cache, value_cache = cache
+          k = MLX::Core.concatenate([key_cache, k], 2)
+          v = MLX::Core.concatenate([value_cache, v], 2)
+        end
+        next_cache = [k, v]
+
+        k_for_attn = __dsl_repeat_kv(k)
+        v_for_attn = __dsl_repeat_kv(v)
+        out = __dsl_attention(q, k_for_attn, v_for_attn, mask)
+        out = MLX::Core.transpose(out, [0, 2, 1, 3])
+        out = MLX::Core.reshape(out, [batch_size, q_len, @num_heads * @head_dim])
+        out = out_proj.call(out)
+        out = MLX::Core.squeeze(out, 0) if q_was_2d
+
+        if @cache_enabled || !cache.nil?
+          [out, next_cache]
+        else
+          out
+        end
+      end
+
+      private
+
+      def __dsl_build_rope(config)
+        return nil if config.nil?
+
+        opts = config.transform_keys(&:to_sym)
+        rope_kwargs = {
+          traditional: opts.fetch(:traditional, false),
+          base: opts.fetch(:base, 10_000.0)
+        }
+        rope_kwargs[:scale] = opts[:scale] if opts.key?(:scale)
+        MLX::NN::RoPE.new(@head_dim, **rope_kwargs)
+      end
+
+      def __dsl_pack_heads(x, heads)
+        batch, length, = x.shape
+        x = MLX::Core.reshape(x, [batch, length, heads, @head_dim])
+        MLX::Core.transpose(x, [0, 2, 1, 3])
+      end
+
+      def __dsl_repeat_kv(x)
+        return x if @kv_repeats == 1
+
+        batch, _heads, length, dim = x.shape
+        expanded = MLX::Core.expand_dims(x, 2)
+        repeated = MLX::Core.concatenate(Array.new(@kv_repeats, expanded), 2)
+        MLX::Core.reshape(repeated, [batch, @num_heads, length, dim])
+      end
+
+      def __dsl_attention(q, k, v, mask)
+        if @backend == :sdpa && MLX::Core.respond_to?(:scaled_dot_product_attention)
+          return MLX::Core.scaled_dot_product_attention(q, k, v, @scale, mask)
+        end
+
+        scores = MLX::Core.matmul(
+          MLX::Core.multiply(q, @scale),
+          MLX::Core.transpose(k, [0, 1, 3, 2])
+        )
+        scores = MLX::Core.add(scores, mask.astype(scores.dtype)) unless mask.nil?
+        probs = MLX::Core.softmax(scores.astype(MLX::Core.float32), -1).astype(scores.dtype)
+        MLX::Core.matmul(probs, v)
+      end
+    end
+  end
+end

data/lib/mlx/dsl/builder.rb

@@ -266,6 +266,14 @@ module MLX
         push(MLX::NN::Transformer.new(*args, **kwargs))
       end
 
+      def attention(*args, **kwargs)
+        push(MLX::DSL::Attention.new(*args, **kwargs))
+      end
+
+      def transformer_block(*args, **kwargs)
+        push(MLX::DSL::TransformerBlock.new(*args, **kwargs))
+      end
+
       def rope(*args, **kwargs)
         push(MLX::NN::RoPE.new(*args, **kwargs))
       end

data/lib/mlx/dsl/config_schema.rb

@@ -0,0 +1,133 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    module ConfigSchema
+      UNSET = Object.new.freeze
+      class DefaultContext
+        def initialize(values)
+          @values = values
+        end
+
+        def method_missing(name, *args, &block)
+          if args.empty? && block.nil? && @values.key?(name.to_sym)
+            return @values[name.to_sym]
+          end
+
+          super
+        end
+
+        def respond_to_missing?(name, include_private = false)
+          @values.key?(name.to_sym) || super
+        end
+      end
+
+      def self.included(base)
+        base.extend(ClassMethods)
+      end
+
+      module ClassMethods
+        def field(name, type = nil, required: false, default: UNSET, &validator)
+          key = name.to_sym
+          config_schema_fields[key] = {
+            type: type,
+            required: !!required,
+            default: default,
+            validator: validator
+          }
+          attr_accessor key unless method_defined?(key) && method_defined?(:"#{key}=")
+        end
+
+        def config_schema_fields
+          @config_schema_fields ||= {}
+        end
+
+        def inherited(subclass)
+          super
+          copied = config_schema_fields.each_with_object({}) do |(key, value), out|
+            out[key] = value.dup
+          end
+          subclass.instance_variable_set(:@config_schema_fields, copied)
+        end
+
+        def from_hash(raw)
+          source = (raw || {}).each_with_object({}) do |(key, value), out|
+            out[key.to_sym] = value
+          end
+          new(**source)
+        end
+
+        private
+
+        def __dsl_call_default(default, resolved)
+          context = DefaultContext.new(resolved)
+          return default unless default.respond_to?(:call)
+          return default.call(context) if default.is_a?(Proc) && default.arity == 1
+          return default.call if !default.is_a?(Proc)
+          return default.call if default.arity.zero?
+
+          default.call(context)
+        end
+
+        def __dsl_validate_field(name, value, spec)
+          type = spec.fetch(:type)
+          if !type.nil? && !value.nil?
+            allowed_types = type.is_a?(Array) ? type : [type]
+            unless allowed_types.any? { |klass| value.is_a?(klass) }
+              raise TypeError,
+                    "config field #{name} must be #{allowed_types.map(&:to_s).join(' or ')}, got #{value.class}"
+            end
+          end
+
+          validator = spec.fetch(:validator)
+          unless validator.nil?
+            if validator.arity == 2
+              validator.call(value, name)
+            else
+              validator.call(value)
+            end
+          end
+
+          value
+        end
+      end
+
+      def initialize(**kwargs)
+        source = kwargs.each_with_object({}) do |(key, value), out|
+          out[key.to_sym] = value
+        end
+        resolved = {}
+        unknown = source.keys - self.class.config_schema_fields.keys
+        unless unknown.empty?
+          names = unknown.map(&:to_s).sort.join(", ")
+          raise ArgumentError, "unknown config field(s): #{names}"
+        end
+
+        self.class.config_schema_fields.each do |name, spec|
+          if source.key?(name)
+            value = source.fetch(name)
+          else
+            default = spec.fetch(:default)
+            if default.equal?(UNSET)
+              if spec.fetch(:required)
+                raise ArgumentError, "missing required config field: #{name}"
+              end
+              next
+            end
+            value = self.class.send(:__dsl_call_default, default, resolved)
+          end
+
+          value = self.class.send(:__dsl_validate_field, name, value, spec)
+          resolved[name] = value
+          public_send(:"#{name}=", value)
+        end
+      end
+
+      def to_h
+        self.class.config_schema_fields.keys.each_with_object({}) do |name, out|
+          out[name.to_s] = public_send(name)
+        end
+      end
+    end
+  end
+end

data/lib/mlx/dsl/generate.rb

@@ -0,0 +1,193 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class Generate
+      def initialize(
+        model:,
+        tokenizer: nil,
+        eos_id: nil,
+        sampler: nil,
+        mode: :decoder_only,
+        decoder_start_id: nil
+      )
+        @model = model
+        @tokenizer = tokenizer
+        @eos_id = eos_id
+        @sampler = { strategy: :argmax }.merge((sampler || {}).transform_keys(&:to_sym))
+        @mode = mode.to_sym
+        @decoder_start_id = decoder_start_id
+      end
+
+      def each_token(prompt: nil, input_ids: nil, max_tokens: 128, **kwargs)
+        return enum_for(__method__, prompt: prompt, input_ids: input_ids, max_tokens: max_tokens, **kwargs) unless block_given?
+
+        case @mode
+        when :decoder_only
+          __dsl_each_decoder_only(prompt: prompt, input_ids: input_ids, max_tokens: max_tokens, **kwargs) do |id, chunk|
+            yield id, chunk
+          end
+        when :encoder_decoder
+          __dsl_each_encoder_decoder(prompt: prompt, input_ids: input_ids, max_tokens: max_tokens, **kwargs) do |id, chunk|
+            yield id, chunk
+          end
+        else
+          raise ArgumentError, "unsupported generation mode: #{@mode.inspect}"
+        end
+
+        self
+      end
+
+      private
+
+      def __dsl_each_decoder_only(prompt:, input_ids:, max_tokens:, **kwargs)
+        tokens = input_ids.nil? ? __dsl_encode(prompt) : input_ids
+        model_input = __dsl_input_array(tokens)
+        logits, cache = __dsl_decode_step(model_input, cache: nil, **kwargs)
+
+        max_tokens.to_i.times do
+          token = __dsl_sample(__dsl_last_logits(logits))
+          token_id = __dsl_token_id(token)
+          chunk = __dsl_decode_token(token_id)
+          yield token_id, chunk
+          break if !@eos_id.nil? && token_id == @eos_id
+
+          next_input = MLX::Core.array([[token_id]], MLX::Core.int32)
+          logits, cache = __dsl_decode_step(next_input, cache: cache, **kwargs)
+        end
+      end
+
+      def __dsl_each_encoder_decoder(prompt:, input_ids:, max_tokens:, **kwargs)
+        tokens = input_ids.nil? ? __dsl_encode(prompt) : input_ids
+        source = __dsl_input_array(tokens)
+
+        if @model.respond_to?(:encode) && @model.respond_to?(:decode)
+          memory = @model.encode(source)
+          start_id = __dsl_decoder_start_id
+          decoder_input = MLX::Core.array([[start_id]], MLX::Core.int32)
+          cache = nil
+
+          max_tokens.to_i.times do
+            decoded = @model.decode(decoder_input, memory, cache: cache, **kwargs)
+            logits, cache = __dsl_split_logits_and_cache(decoded, cache)
+            token = __dsl_sample(__dsl_last_logits(logits))
+            token_id = __dsl_token_id(token)
+            chunk = __dsl_decode_token(token_id)
+            yield token_id, chunk
+            break if !@eos_id.nil? && token_id == @eos_id
+
+            decoder_input = MLX::Core.array([[token_id]], MLX::Core.int32)
+          end
+          return
+        end
+
+        # Fallback path for model.call-style APIs.
+        __dsl_each_decoder_only(prompt: prompt, input_ids: tokens, max_tokens: max_tokens, **kwargs) do |id, chunk|
+          yield id, chunk
+        end
+      end
+
+      def __dsl_decode_step(input_ids, cache:, **kwargs)
+        output = @model.call(input_ids, cache: cache, **kwargs)
+        __dsl_split_logits_and_cache(output, cache)
+      end
+
+      def __dsl_split_logits_and_cache(output, fallback_cache)
+        if output.is_a?(Array) && output.length == 2
+          [output[0], output[1]]
+        else
+          [output, fallback_cache]
+        end
+      end
+
+      def __dsl_last_logits(logits)
+        return logits if logits.ndim == 2
+        return logits if logits.ndim == 1
+
+        index = MLX::Core.array([logits.shape[1] - 1], MLX::Core.int32)
+        MLX::Core.squeeze(MLX::Core.take(logits, index, 1), 1)
+      end
+
+      def __dsl_sample(logits)
+        strategy = @sampler.fetch(:strategy, :argmax).to_sym
+        temperature = @sampler.fetch(:temperature, 1.0).to_f
+
+        return MLX::Core.argmax(logits, -1) if strategy == :argmax || temperature.zero?
+
+        case strategy
+        when :top_k
+          __dsl_top_k_sample(logits, k: Integer(@sampler.fetch(:k, 40)), temperature: temperature)
+        when :temperature, :categorical
+          __dsl_temperature_sample(logits, temperature: temperature)
+        else
+          raise ArgumentError, "unsupported sampler strategy: #{strategy.inspect}"
+        end
+      end
+
+      def __dsl_temperature_sample(logits, temperature:)
+        scaled = if temperature == 1.0
+          logits
+        else
+          MLX::Core.multiply(logits, 1.0 / temperature)
+        end
+        MLX::Core.categorical(scaled)
+      end
+
+      def __dsl_top_k_sample(logits, k:, temperature:)
+        rows = logits.ndim == 1 ? [logits.to_a] : logits.to_a
+        masked = rows.map do |row|
+          pairs = row.each_with_index.sort_by { |(value, _index)| -value }
+          keep = pairs.first([k, row.length].min).map(&:last)
+          filtered = Array.new(row.length, -Float::INFINITY)
+          keep.each { |idx| filtered[idx] = row[idx] }
+          filtered
+        end
+        masked_logits = MLX::Core.array(masked, logits.dtype)
+        __dsl_temperature_sample(masked_logits, temperature: temperature)
+      end
+
+      def __dsl_encode(prompt)
+        raise ArgumentError, "prompt/input_ids required when tokenizer is unavailable" if @tokenizer.nil?
+
+        @tokenizer.encode(prompt.to_s)
+      end
+
+      def __dsl_input_array(tokens)
+        if tokens.is_a?(MLX::Core::Array)
+          return tokens if tokens.ndim > 1
+
+          return MLX::Core.expand_dims(tokens.astype(MLX::Core.int32), 0)
+        end
+
+        arr = tokens.to_a
+        nested = arr.empty? ? [[]] : (arr.first.is_a?(Array) ? arr : [arr])
+        MLX::Core.array(nested, MLX::Core.int32)
+      end
+
+      def __dsl_token_id(token)
+        return token.item.to_i if token.respond_to?(:item)
+
+        value = token.to_a
+        if value.is_a?(Array)
+          first = value.first
+          return first.first.to_i if first.is_a?(Array)
+          return first.to_i
+        end
+        value.to_i
+      end
+
+      def __dsl_decode_token(token_id)
+        return nil if @tokenizer.nil? || !@tokenizer.respond_to?(:decode)
+
+        @tokenizer.decode([token_id])
+      end
+
+      def __dsl_decoder_start_id
+        return @decoder_start_id unless @decoder_start_id.nil?
+        return @tokenizer.decoder_start_id if !@tokenizer.nil? && @tokenizer.respond_to?(:decoder_start_id)
+
+        raise ArgumentError, "decoder_start_id is required for encoder-decoder mode"
+      end
+    end
+  end
+end

data/lib/mlx/dsl/kv_cache.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class KVCache
+      attr_reader :num_layers
+
+      def initialize(num_layers:)
+        @num_layers = Integer(num_layers)
+        raise ArgumentError, "num_layers must be non-negative" if @num_layers.negative?
+
+        @layers = Array.new(@num_layers)
+      end
+
+      def layer(index)
+        @layers.fetch(__dsl_index(index))
+      end
+
+      def []=(index, value)
+        @layers[__dsl_index(index)] = value
+      end
+
+      def offset(layer:)
+        state = self.layer(layer)
+        return 0 if state.nil?
+
+        keys, = state
+        keys.shape[2]
+      end
+
+      def append(layer:, keys:, values:)
+        idx = __dsl_index(layer)
+        current = @layers[idx]
+        if current.nil?
+          @layers[idx] = [keys, values]
+          return @layers[idx]
+        end
+
+        key_cache, value_cache = current
+        next_keys = MLX::Core.concatenate([key_cache, keys], 2)
+        next_values = MLX::Core.concatenate([value_cache, values], 2)
+        @layers[idx] = [next_keys, next_values]
+      end
+
+      def truncate!(tokens:, layer: nil)
+        keep = Integer(tokens)
+        if layer.nil?
+          @layers.each_index { |idx| __dsl_truncate_layer!(idx, keep) }
+        else
+          __dsl_truncate_layer!(__dsl_index(layer), keep)
+        end
+        self
+      end
+
+      def reset!(layer: nil)
+        if layer.nil?
+          @layers.map! { nil }
+        else
+          @layers[__dsl_index(layer)] = nil
+        end
+        self
+      end
+
+      private
+
+      def __dsl_index(index)
+        idx = Integer(index)
+        if idx.negative? || idx >= @num_layers
+          raise IndexError, "layer index #{idx} out of range (0...#{@num_layers})"
+        end
+
+        idx
+      end
+
+      def __dsl_truncate_layer!(idx, keep)
+        state = @layers[idx]
+        return if state.nil?
+
+        if keep <= 0
+          @layers[idx] = nil
+          return
+        end
+
+        keys, values = state
+        total = keys.shape[2]
+        return if keep >= total
+
+        start = total - keep
+        indices = MLX::Core.arange(start, total, 1, MLX::Core.int32)
+        trimmed_keys = MLX::Core.take(keys, indices, 2)
+        trimmed_values = MLX::Core.take(values, indices, 2)
+        @layers[idx] = [trimmed_keys, trimmed_values]
+      end
+    end
+  end
+end

data/lib/mlx/dsl/masks.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    module Masks
+      module_function
+
+      def causal(length:, offset: 0, dtype: MLX::Core.float32)
+        length = Integer(length)
+        offset = Integer(offset)
+        raise ArgumentError, "length must be non-negative" if length.negative?
+        raise ArgumentError, "offset must be non-negative" if offset.negative?
+
+        rinds = MLX::Core.arange(0, offset + length, 1)
+        linds = if offset.zero?
+          rinds
+        else
+          MLX::Core.arange(offset, offset + length, 1)
+        end
+        lhs = MLX::Core.expand_dims(linds, 1)
+        rhs = MLX::Core.expand_dims(rinds, 0)
+        mask = MLX::Core.less(lhs, rhs).astype(dtype)
+        min_value = if MLX::Core.respond_to?(:finfo)
+          MLX::Core.finfo(dtype).min
+        else
+          -1e9
+        end
+        MLX::Core.multiply(mask, min_value)
+      end
+    end
+  end
+end

data/lib/mlx/dsl/positions.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    module Positions
+      module_function
+
+      def ids_like(input_ids, offset: 0, dtype: nil)
+        shape = input_ids.shape
+        seq_len = shape[-1]
+        dtype ||= input_ids.respond_to?(:dtype) ? input_ids.dtype : MLX::Core.int32
+
+        base = MLX::Core.arange(offset.to_i, offset.to_i + seq_len, 1, dtype)
+        return base if shape.length == 1
+
+        reshape_dims = Array.new(shape.length, 1)
+        reshape_dims[-1] = seq_len
+        expanded = MLX::Core.reshape(base, reshape_dims)
+        MLX::Core.broadcast_to(expanded, shape)
+      end
+
+      def offset_from_cache(cache, layer: 0)
+        return 0 if cache.nil?
+        return cache.offset(layer: layer) if cache.respond_to?(:offset)
+
+        if cache.respond_to?(:[]) && !cache[layer].nil?
+          keys, = cache[layer]
+          return keys.shape[2]
+        end
+
+        0
+      end
+    end
+  end
+end

data/lib/mlx/dsl/run_stack.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    def self.run_stack(layers, input, cache: nil, **kwargs)
+      modules = layers.to_a
+      if cache.is_a?(MLX::DSL::KVCache)
+        hidden = input
+        modules.each_with_index do |layer, index|
+          hidden, next_cache = __dsl_run_stack_layer(
+            layer,
+            hidden,
+            kwargs,
+            cache: cache.layer(index),
+            use_cache: true
+          )
+          cache[index] = next_cache
+        end
+        return [hidden, cache]
+      end
+
+      use_cache = !cache.nil?
+      cache_state = if use_cache
+        entries = cache.to_a
+        entries.length < modules.length ? entries + Array.new(modules.length - entries.length) : entries.dup
+      else
+        nil
+      end
+
+      hidden = input
+      modules.each_with_index do |layer, index|
+        layer_cache = use_cache ? cache_state[index] : nil
+        hidden, next_cache = __dsl_run_stack_layer(
+          layer,
+          hidden,
+          kwargs,
+          cache: layer_cache,
+          use_cache: use_cache
+        )
+        cache_state[index] = next_cache if use_cache
+      end
+
+      use_cache ? [hidden, cache_state] : hidden
+    end
+
+    def self.__dsl_run_stack_layer(layer, hidden, kwargs, cache:, use_cache:)
+      call_kwargs = kwargs.dup
+      call_kwargs[:cache] = cache if use_cache
+      result = layer.call(hidden, **call_kwargs)
+
+      if use_cache && result.is_a?(Array) && result.length == 2
+        [result[0], result[1]]
+      else
+        [result, cache]
+      end
+    rescue ArgumentError => e
+      if use_cache && e.message.include?("unknown keyword: :cache")
+        result = layer.call(hidden, **kwargs)
+        if result.is_a?(Array) && result.length == 2
+          return [result[0], result[1]]
+        end
+        return [result, cache]
+      end
+      raise
+    end
+    private_class_method :__dsl_run_stack_layer
+  end
+end