mlx 0.30.7 → 0.30.7.2

data/lib/mlx/dsl/data_pipeline.rb

@@ -0,0 +1,284 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    module Data
+      def self.from(source = nil, &block)
+        if !source.nil? && block_given?
+          raise ArgumentError, "data pipeline source accepts either a source or block, not both"
+        end
+
+        producer = block_given? ? block : source
+        if producer.nil?
+          raise ArgumentError, "data pipeline requires a source enumerable or source block"
+        end
+
+        Pipeline.new(__dsl_factory_for(producer))
+      end
+
+      def self.pipeline(source = nil, &block)
+        from(source, &block)
+      end
+
+      def self.__dsl_factory_for(producer)
+        if producer.respond_to?(:call)
+          lambda do
+            __dsl_to_enumerator(producer.call)
+          end
+        else
+          lambda do
+            if producer.respond_to?(:rewind)
+              begin
+                producer.rewind
+              rescue StandardError
+                # Keep default Enumerable semantics when rewind is unavailable at runtime.
+              end
+            end
+            __dsl_to_enumerator(producer)
+          end
+        end
+      end
+      private_class_method :__dsl_factory_for
+
+      def self.__dsl_to_enumerator(value)
+        unless value.respond_to?(:each)
+          raise ArgumentError, "data pipeline source must respond to #each"
+        end
+
+        value.to_enum
+      end
+      private_class_method :__dsl_to_enumerator
+
+      class Pipeline
+        include Enumerable
+
+        def initialize(factory)
+          @factory = factory
+        end
+
+        def each
+          enum = @factory.call
+          return enum unless block_given?
+
+          enum.each { |item| yield item }
+        end
+
+        def map(&block)
+          raise ArgumentError, "pipeline map requires a block" unless block_given?
+
+          self.class.new(lambda {
+            upstream = @factory.call
+            Enumerator.new do |y|
+              index = 0
+              upstream.each do |item|
+                y << __dsl_call_with_context(block, item, index, "pipeline map")
+                index += 1
+              end
+            end
+          })
+        end
+
+        def filter(&block)
+          raise ArgumentError, "pipeline filter requires a block" unless block_given?
+
+          self.class.new(lambda {
+            upstream = @factory.call
+            Enumerator.new do |y|
+              index = 0
+              upstream.each do |item|
+                y << item if __dsl_call_with_context(block, item, index, "pipeline filter")
+                index += 1
+              end
+            end
+          })
+        end
+
+        def batch(size, drop_last: false)
+          batch_size = size.to_i
+          raise ArgumentError, "pipeline batch size must be positive" if batch_size <= 0
+
+          self.class.new(lambda {
+            upstream = @factory.call
+            Enumerator.new do |y|
+              chunk = []
+              upstream.each do |item|
+                chunk << item
+                if chunk.length == batch_size
+                  y << chunk
+                  chunk = []
+                end
+              end
+              y << chunk unless drop_last || chunk.empty?
+            end
+          })
+        end
+
+        def take(count)
+          limit = count.to_i
+          raise ArgumentError, "pipeline take count must be non-negative" if limit.negative?
+
+          self.class.new(lambda {
+            upstream = @factory.call
+            Enumerator.new do |y|
+              seen = 0
+              while seen < limit
+                begin
+                  y << upstream.next
+                  seen += 1
+                rescue StopIteration
+                  break
+                end
+              end
+            end
+          })
+        end
+
+        def repeat(times = nil)
+          if times.nil?
+            self.class.new(lambda {
+              Enumerator.new do |y|
+                loop do
+                  upstream = @factory.call
+                  produced = false
+                  upstream.each do |item|
+                    produced = true
+                    y << item
+                  end
+                  break unless produced
+                end
+              end
+            })
+          else
+            cycles = times.to_i
+            raise ArgumentError, "pipeline repeat count must be non-negative" if cycles.negative?
+
+            self.class.new(lambda {
+              Enumerator.new do |y|
+                cycles.times do
+                  @factory.call.each do |item|
+                    y << item
+                  end
+                end
+              end
+            })
+          end
+        end
+
+        def shuffle(seed: nil, random: nil)
+          if !seed.nil? && !random.nil?
+            raise ArgumentError, "pipeline shuffle accepts either seed: or random:, not both"
+          end
+
+          self.class.new(lambda {
+            items = @factory.call.to_a
+            rng = if !random.nil?
+              random
+            elsif !seed.nil?
+              Random.new(seed.to_i)
+            else
+              Random.new
+            end
+            items.shuffle(random: rng).to_enum
+          })
+        end
+
+        def prefetch(size = 1)
+          prefetch_size = size.to_i
+          raise ArgumentError, "pipeline prefetch size must be positive" if prefetch_size <= 0
+
+          self.class.new(lambda {
+            upstream = @factory.call
+            Enumerator.new do |y|
+              buffer = []
+
+              prefetch_size.times do
+                begin
+                  buffer << upstream.next
+                rescue StopIteration
+                  break
+                end
+              end
+
+              until buffer.empty?
+                y << buffer.shift
+                begin
+                  buffer << upstream.next
+                rescue StopIteration
+                  # Exhausted upstream; continue draining buffer.
+                end
+              end
+            end
+          })
+        end
+
+        private
+
+        def __dsl_call_with_context(callable, item, index, label)
+          values = {
+            item: item,
+            index: index,
+            pipeline: self
+          }
+          return callable.call(item, index) unless callable.respond_to?(:parameters)
+
+          params = callable.parameters
+          return callable.call(item, index) if params.empty?
+
+          args = __dsl_build_positional_args(
+            params,
+            values,
+            [[:item, item], [:index, index], [:pipeline, self]],
+            label
+          )
+          kwargs = __dsl_build_keyword_args(params, values, label)
+          return callable.call(*args) if kwargs.empty?
+
+          callable.call(*args, **kwargs)
+        end
+
+        def __dsl_build_positional_args(params, values, fallback_pairs, label)
+          queue = fallback_pairs.dup
+          args = []
+          params.each do |type, name|
+            next unless type == :req || type == :opt
+
+            if !name.nil? && values.key?(name)
+              args << values.fetch(name)
+              queue.reject! { |key, _value| key == name }
+              next
+            end
+
+            if queue.empty?
+              raise ArgumentError, "#{label} has unsupported required positional argument: #{name.inspect}" if type == :req
+              break
+            end
+
+            _key, value = queue.shift
+            args << value
+          end
+          args
+        end
+
+        def __dsl_build_keyword_args(params, values, label)
+          return values.dup if params.any? { |type, _name| type == :keyrest }
+
+          required_keys = params.each_with_object([]) do |(type, name), out|
+            out << name if type == :keyreq
+          end
+          missing = required_keys.reject { |name| values.key?(name) }
+          unless missing.empty?
+            raise ArgumentError, "#{label} requires unsupported keyword argument(s): #{missing.map(&:inspect).join(", ")}"
+          end
+
+          accepted_keys = params.each_with_object([]) do |(type, name), out|
+            out << name if type == :key || type == :keyreq
+          end
+
+          values.each_with_object({}) do |(name, value), out|
+            out[name] = value if accepted_keys.include?(name)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/mlx/dsl/experiment.rb

@@ -0,0 +1,154 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    def self.experiment(name = nil, &block)
+      instance = Experiment.new(name: name)
+      instance.instance_eval(&block) if block_given?
+      instance
+    end
+
+    class Experiment
+      attr_reader :name
+
+      def initialize(name: nil)
+        @name = name
+        @model_source = nil
+        @optimizer_source = nil
+        @trainer_source = nil
+        @trainer_kwargs = {}
+        @loss_block = nil
+        @data_config = { train: nil, validation: nil, fit: {} }
+        @artifact_config = {}
+        @last_trainer = nil
+        @last_report = nil
+      end
+
+      def model(value = nil, &block)
+        if !value.nil? && block_given?
+          raise ArgumentError, "model accepts either a value argument or block, not both"
+        end
+
+        @model_source = block_given? ? block : value
+        self
+      end
+
+      def optimizer(value = nil, &block)
+        if !value.nil? && block_given?
+          raise ArgumentError, "optimizer accepts either a value argument or block, not both"
+        end
+
+        @optimizer_source = block_given? ? block : value
+        self
+      end
+
+      def trainer(value = nil, **kwargs, &block)
+        if value.is_a?(MLX::DSL::Trainer)
+          if !kwargs.empty? || block_given?
+            raise ArgumentError, "trainer instance injection cannot be combined with trainer kwargs or loss block"
+          end
+          @trainer_source = value
+          return self
+        end
+
+        unless value.nil?
+          raise ArgumentError, "trainer positional argument must be an MLX::DSL::Trainer instance"
+        end
+
+        @trainer_source = nil
+        @trainer_kwargs = kwargs.dup
+        @loss_block = block if block_given?
+        self
+      end
+
+      def data(train: nil, validation: :__dsl_unset__, **fit_kwargs)
+        @data_config[:train] = train unless train.nil?
+        @data_config[:validation] = validation unless validation == :__dsl_unset__
+        @data_config[:fit].merge!(fit_kwargs)
+        self
+      end
+
+      def artifacts(**kwargs)
+        @artifact_config.merge!(kwargs)
+        self
+      end
+
+      def run(report: false, **overrides)
+        dataset, fit_kwargs = __dsl_resolve_fit_call(overrides)
+        active_trainer = __dsl_resolve_trainer
+        result = if report
+          active_trainer.fit_report(dataset, **fit_kwargs)
+        else
+          active_trainer.fit(dataset, **fit_kwargs)
+        end
+        @last_report = result if report
+        result
+      end
+
+      def report(**overrides)
+        run(report: true, **overrides)
+      end
+
+      def save_run_bundle(path, report: nil, config: {}, **overrides)
+        active_report = report
+        if active_report.nil?
+          active_report = if !@last_report.nil?
+            @last_report
+          else
+            self.report(**overrides)
+          end
+        end
+
+        __dsl_resolve_trainer.save_run_bundle(path, report: active_report, config: config)
+      end
+
+      private
+
+      def __dsl_resolve_fit_call(overrides)
+        fit_kwargs = @data_config.fetch(:fit).dup
+        fit_kwargs.merge!(@artifact_config)
+        fit_kwargs[:validation_data] = @data_config[:validation] if !@data_config[:validation].nil? && !fit_kwargs.key?(:validation_data)
+
+        incoming = overrides.dup
+        dataset = if incoming.key?(:dataset)
+          incoming.delete(:dataset)
+        elsif incoming.key?(:train)
+          incoming.delete(:train)
+        else
+          @data_config[:train]
+        end
+        if dataset.nil?
+          raise ArgumentError, "experiment run requires a train dataset via data(train:) or run(dataset:)"
+        end
+
+        [dataset, fit_kwargs.merge(incoming)]
+      end
+
+      def __dsl_resolve_trainer
+        return @trainer_source if @trainer_source.is_a?(MLX::DSL::Trainer)
+        return @last_trainer unless @last_trainer.nil?
+
+        model = __dsl_resolve_source(@model_source, "model")
+        optimizer = __dsl_resolve_source(@optimizer_source, "optimizer")
+        unless model.respond_to?(:trainer)
+          raise ArgumentError, "experiment model must respond to #trainer when trainer instance is not injected"
+        end
+        unless @loss_block.respond_to?(:call)
+          raise ArgumentError, "experiment trainer requires a loss block when trainer instance is not injected"
+        end
+
+        @last_trainer = model.trainer(optimizer: optimizer, **@trainer_kwargs, &@loss_block)
+      end
+
+      def __dsl_resolve_source(source, label)
+        value = source
+        value = value.call if value.respond_to?(:call)
+        if value.nil?
+          raise ArgumentError, "experiment #{label} section is required when trainer instance is not injected"
+        end
+
+        value
+      end
+    end
+  end
+end

data/lib/mlx/dsl/graph_modules.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class Callable < MLX::NN::Module
+      def initialize(callable = nil, &block)
+        super()
+        if !callable.nil? && block_given?
+          raise ArgumentError, "callable layer accepts either a callable argument or block, not both"
+        end
+
+        @callable = callable.nil? ? block : callable
+        unless @callable.respond_to?(:call)
+          raise ArgumentError, "callable layer requires a callable argument or block"
+        end
+      end
+
+      def call(*args, **kwargs)
+        return @callable.call(*args) if kwargs.empty?
+
+        @callable.call(*args, **kwargs)
+      end
+    end
+
+    class Residual < MLX::NN::Module
+      def initialize(module_obj)
+        super()
+        self.module_obj = module_obj
+      end
+
+      def call(*args, **kwargs)
+        raise ArgumentError, "residual module expects at least one positional input" if args.empty?
+
+        identity = args[0]
+        transformed = module_obj.call(*args, **kwargs)
+        MLX::Core.add(identity, transformed)
+      end
+    end
+
+    class Parallel < MLX::NN::Module
+      def initialize(*modules)
+        super()
+        self.layers = modules
+      end
+
+      def call(*args, **kwargs)
+        layers.map do |layer|
+          layer.call(*args, **kwargs)
+        end
+      end
+    end
+
+    class Concat < MLX::NN::Module
+      def initialize(*modules, axis: -1)
+        super()
+        self.layers = modules
+        @axis = axis
+      end
+
+      def call(*args, **kwargs)
+        outputs = layers.map do |layer|
+          layer.call(*args, **kwargs)
+        end
+        MLX::Core.concatenate(outputs, @axis)
+      end
+    end
+
+    class Reduce < MLX::NN::Module
+      def initialize(*modules, mode: :sum)
+        super()
+        self.layers = modules
+        @mode = mode.to_sym
+      end
+
+      def call(*args, **kwargs)
+        outputs = layers.map do |layer|
+          layer.call(*args, **kwargs)
+        end
+
+        case @mode
+        when :sum
+          outputs.reduce do |acc, item|
+            MLX::Core.add(acc, item)
+          end
+        else
+          raise ArgumentError, "unsupported reduce mode: #{@mode.inspect}"
+        end
+      end
+    end
+  end
+end

data/lib/mlx/dsl/model.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class Model < MLX::NN::Module
+      include ModelMixin
+    end
+  end
+end