mlx 0.30.7 → 0.30.7.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: ce770b59d71e1d5fbf8697bfebde05006fa52a770f8526575f94944cf161d05b
-  data.tar.gz: 814aa4e6c063b3b36c3b1d3c4c26d8f5fa3952a66947d294a8edb3864b07ece2
+  metadata.gz: 215a912d2353fd5edaa60e320a5b857aa13009e55f3a190b21b2ffe5735f37af
+  data.tar.gz: 1c9b4279f8077e3cd067354ea692e1b24248afe7262d88d4569c566c52f5a158
 SHA512:
-  metadata.gz: 3702f6d4445ea4af978ffcebb71dcef91f8e28d4f4d79a1d16a49b66848023a84b7b8ff2a614fa7b2c236d324dfb92b8ccb01e70ff59c53234ac28ee6fc09b39
-  data.tar.gz: 52bf528fee068f422dac611fdca22fae8089a7ea9a3a4d0cd42a21aecc53b5219b0dcf83d79f5fa907a1c99f97acb831d952474963071b566b3483e91c4a4a72
+  metadata.gz: 66abcbd58ccfc04186df11b0d2b6445c7d1e0ab4a36451742755f6fcf41022363403536c3d27640935364c58231c4c4a03e39ceba97617a59f2ad69acf23dc16
+  data.tar.gz: ba7ad07ccd31e94bdf3fdee73117f069c6ee22c11cdfc3f2470eedee9ee0bc9e976970bab21f156c003b658723a0e3fb2f63f62d0e11d33de3a61fc1d9121711

data/ext/mlx/native.cpp

@@ -7884,9 +7884,6 @@ extern "C" void Init_native(void) {
       "scaled_dot_product_attention",
       RUBY_METHOD_FUNC(core_scaled_dot_product_attention),
       -1);
-  rb_define_singleton_method(
-      mCore, "scaled_dot_product_attention", RUBY_METHOD_FUNC(core_scaled_dot_product_attention), -1);
-  rb_define_singleton_method(mCore, "scaled_dot_product_attention", RUBY_METHOD_FUNC(core_scaled_dot_product_attention), -1);
   rb_define_singleton_method(mCore, "arange", RUBY_METHOD_FUNC(core_arange), -1);
   rb_define_singleton_method(mCore, "linspace", RUBY_METHOD_FUNC(core_linspace), -1);
   rb_define_singleton_method(mCore, "zeros", RUBY_METHOD_FUNC(core_zeros), -1);
@@ -8023,5 +8020,4 @@ extern "C" void Init_native(void) {
       "precompiled_cuda_kernel",
       RUBY_METHOD_FUNC(core_precompiled_cuda_kernel),
       -1);
-  rb_define_singleton_method(mCore, "precompiled_cuda_kernel", RUBY_METHOD_FUNC(core_precompiled_cuda_kernel), -1);
 }

data/lib/mlx/core.rb

@@ -335,6 +335,12 @@ module MLX
       alias_method :native_export_to_dot,
                    :export_to_dot if method_defined?(:export_to_dot) && !method_defined?(:native_export_to_dot)
 
+      %i[savez savez_compressed].each do |method_name|
+        if method_defined?(method_name) && instance_method(method_name).owner == self
+          remove_method(method_name)
+        end
+      end
+
       ARRAY_LEAF = :__mlx_array_leaf__
 
       def load(file, format = nil, return_metadata = false)
@@ -963,7 +969,8 @@ module MLX
         end
       end
 
-      alias eql? ==
+      remove_method(:eql?) if method_defined?(:eql?) && instance_method(:eql?).owner == self
+      alias_method :eql?, :==
     end
 
     class Array

data/lib/mlx/distributed_utils/launch.rb

@@ -373,12 +373,18 @@ module MLX
         opts[:env] = hostfile.envs + opts[:env]
 
         command = rest.dup
-        script = Pathname.new(command.first)
-        if script.file?
+        command_name = command.first.to_s
+        script = Pathname.new(command_name)
+        explicit_path = command_name.include?(File::SEPARATOR) || command_name.start_with?(".", "~")
+
+        if explicit_path && script.file?
           command[0] = opts[:python]
           command.insert(1, script.realpath.to_s)
-        elsif (resolved = find_executable(command.first))
+        elsif (resolved = find_executable(command_name))
           command[0] = resolved
+        elsif script.file?
+          command[0] = opts[:python]
+          command.insert(1, script.realpath.to_s)
         elsif opts[:verify_script]
           raise ArgumentError, "Invalid script or command #{command.first}"
         end

data/lib/mlx/dsl/builder.rb

@@ -0,0 +1,377 @@
+# frozen_string_literal: true
+
+module MLX
+  module DSL
+    class Builder
+      def initialize(owner = nil)
+        @owner = owner
+        @collector = nil
+      end
+
+      def build(&block)
+        raise ArgumentError, "builder requires a block" unless block_given?
+
+        instance_eval(&block)
+      end
+
+      def sequential(*modules, &block)
+        collected = __dsl_modules_from(modules, &block)
+        push(MLX::NN::Sequential.new(*collected))
+      end
+
+      def layer(entry = nil, *args, **kwargs, &block)
+        if !entry.nil? && block_given?
+          raise ArgumentError, "layer accepts either a module entry or block, not both"
+        end
+
+        if block_given?
+          return push(MLX::DSL::Callable.new(&block))
+        end
+
+        if entry.nil?
+          raise ArgumentError, "layer requires a module entry or block"
+        end
+
+        if entry.is_a?(MLX::NN::Module)
+          __dsl_reject_layer_constructor_args!(args, kwargs, entry.class)
+          return push(entry)
+        end
+
+        if entry.is_a?(Class)
+          unless entry <= MLX::NN::Module
+            raise TypeError, "layer class must inherit from MLX::NN::Module"
+          end
+          return push(entry.new(*args, **kwargs))
+        end
+
+        if entry.respond_to?(:call)
+          __dsl_reject_layer_constructor_args!(args, kwargs, entry.class)
+          return push(MLX::DSL::Callable.new(entry))
+        end
+
+        raise TypeError, "layer requires an MLX::NN::Module instance, MLX::NN::Module class, callable, or block"
+      end
+
+      def residual(module_obj = nil, &block)
+        modules = __dsl_modules_from(module_obj.nil? ? [] : [module_obj], &block)
+        raise ArgumentError, "residual requires at least one module" if modules.empty?
+
+        target = if modules.length == 1
+          modules[0]
+        else
+          MLX::NN::Sequential.new(*modules)
+        end
+        push(MLX::DSL::Residual.new(target))
+      end
+
+      def branch(*modules, &block)
+        collected = __dsl_modules_from(modules, &block)
+        raise ArgumentError, "branch requires at least one module" if collected.empty?
+
+        push(MLX::DSL::Parallel.new(*collected))
+      end
+
+      def concat(*modules, axis: -1, &block)
+        collected = __dsl_modules_from(modules, &block)
+        raise ArgumentError, "concat requires at least one module" if collected.empty?
+
+        push(MLX::DSL::Concat.new(*collected, axis: axis))
+      end
+
+      def sum(*modules, &block)
+        collected = __dsl_modules_from(modules, &block)
+        raise ArgumentError, "sum requires at least one module" if collected.empty?
+
+        push(MLX::DSL::Reduce.new(*collected, mode: :sum))
+      end
+
+      def fn(callable = nil, &block)
+        push(MLX::DSL::Callable.new(callable, &block))
+      end
+      alias_method :lambda_layer, :fn
+
+      def repeat_layers(count, &block)
+        entries = __dsl_collect_repeated_entries(count, &block)
+        layers = entries.map { |entry| __dsl_normalize_module_entry(entry) }
+        layers.each { |layer| push(layer) }
+        layers
+      end
+
+      def stack(count, layer_class = nil, *args, **kwargs, &block)
+        if !layer_class.nil? && block_given?
+          raise ArgumentError, "stack accepts either a layer class or block, not both"
+        end
+
+        layers = if layer_class.nil?
+          __dsl_collect_repeated_entries(count, &block).map { |entry| __dsl_normalize_module_entry(entry) }
+        else
+          __dsl_build_class_stack_layers(count, layer_class, args, kwargs)
+        end
+        push(MLX::NN::Sequential.new(*layers))
+      end
+
+      def identity(*args, **kwargs)
+        push(MLX::NN::Identity.new(*args, **kwargs))
+      end
+
+      def embedding(*args, **kwargs)
+        push(MLX::NN::Embedding.new(*args, **kwargs))
+      end
+
+      def linear(*args, **kwargs)
+        push(MLX::NN::Linear.new(*args, **kwargs))
+      end
+
+      def bilinear(*args, **kwargs)
+        push(MLX::NN::Bilinear.new(*args, **kwargs))
+      end
+
+      def relu
+        push(MLX::NN::ReLU.new)
+      end
+
+      def relu6
+        push(MLX::NN::ReLU6.new)
+      end
+
+      def leaky_relu(*args)
+        push(MLX::NN::LeakyReLU.new(*args))
+      end
+
+      def gelu(*args, **kwargs)
+        push(MLX::NN::GELU.new(*args, **kwargs))
+      end
+
+      def tanh
+        push(MLX::NN::Tanh.new)
+      end
+
+      def sigmoid
+        push(MLX::NN::Sigmoid.new)
+      end
+
+      def dropout(*args)
+        push(MLX::NN::Dropout.new(*args))
+      end
+
+      def dropout2d(*args)
+        push(MLX::NN::Dropout2d.new(*args))
+      end
+
+      def dropout3d(*args)
+        push(MLX::NN::Dropout3d.new(*args))
+      end
+
+      def conv1d(*args, **kwargs)
+        push(MLX::NN::Conv1d.new(*args, **kwargs))
+      end
+
+      def conv2d(*args, **kwargs)
+        push(MLX::NN::Conv2d.new(*args, **kwargs))
+      end
+
+      def conv3d(*args, **kwargs)
+        push(MLX::NN::Conv3d.new(*args, **kwargs))
+      end
+
+      def conv_transpose1d(*args, **kwargs)
+        push(MLX::NN::ConvTranspose1d.new(*args, **kwargs))
+      end
+
+      def conv_transpose2d(*args, **kwargs)
+        push(MLX::NN::ConvTranspose2d.new(*args, **kwargs))
+      end
+
+      def conv_transpose3d(*args, **kwargs)
+        push(MLX::NN::ConvTranspose3d.new(*args, **kwargs))
+      end
+
+      def layer_norm(*args, **kwargs)
+        push(MLX::NN::LayerNorm.new(*args, **kwargs))
+      end
+
+      def rms_norm(*args, **kwargs)
+        push(MLX::NN::RMSNorm.new(*args, **kwargs))
+      end
+
+      def batch_norm(*args, **kwargs)
+        push(MLX::NN::BatchNorm.new(*args, **kwargs))
+      end
+
+      def instance_norm(*args, **kwargs)
+        push(MLX::NN::InstanceNorm.new(*args, **kwargs))
+      end
+
+      def group_norm(*args, **kwargs)
+        push(MLX::NN::GroupNorm.new(*args, **kwargs))
+      end
+
+      def max_pool2d(*args, **kwargs)
+        push(MLX::NN::MaxPool2d.new(*args, **kwargs))
+      end
+
+      def avg_pool2d(*args, **kwargs)
+        push(MLX::NN::AvgPool2d.new(*args, **kwargs))
+      end
+
+      def max_pool1d(*args, **kwargs)
+        push(MLX::NN::MaxPool1d.new(*args, **kwargs))
+      end
+
+      def avg_pool1d(*args, **kwargs)
+        push(MLX::NN::AvgPool1d.new(*args, **kwargs))
+      end
+
+      def max_pool3d(*args, **kwargs)
+        push(MLX::NN::MaxPool3d.new(*args, **kwargs))
+      end
+
+      def avg_pool3d(*args, **kwargs)
+        push(MLX::NN::AvgPool3d.new(*args, **kwargs))
+      end
+
+      def rnn(*args, **kwargs)
+        push(MLX::NN::RNN.new(*args, **kwargs))
+      end
+
+      def gru(*args, **kwargs)
+        push(MLX::NN::GRU.new(*args, **kwargs))
+      end
+
+      def lstm(*args, **kwargs)
+        push(MLX::NN::LSTM.new(*args, **kwargs))
+      end
+
+      def multi_head_attention(*args, **kwargs)
+        push(MLX::NN::MultiHeadAttention.new(*args, **kwargs))
+      end
+
+      def transformer_encoder_layer(*args, **kwargs)
+        push(MLX::NN::TransformerEncoderLayer.new(*args, **kwargs))
+      end
+
+      def transformer_encoder(*args, **kwargs)
+        push(MLX::NN::TransformerEncoder.new(*args, **kwargs))
+      end
+
+      def transformer_decoder_layer(*args, **kwargs)
+        push(MLX::NN::TransformerDecoderLayer.new(*args, **kwargs))
+      end
+
+      def transformer_decoder(*args, **kwargs)
+        push(MLX::NN::TransformerDecoder.new(*args, **kwargs))
+      end
+
+      def transformer(*args, **kwargs)
+        push(MLX::NN::Transformer.new(*args, **kwargs))
+      end
+
+      def rope(*args, **kwargs)
+        push(MLX::NN::RoPE.new(*args, **kwargs))
+      end
+
+      def sinusoidal_positional_encoding(*args, **kwargs)
+        push(MLX::NN::SinusoidalPositionalEncoding.new(*args, **kwargs))
+      end
+
+      def alibi(*args, **kwargs)
+        push(MLX::NN::ALiBi.new(*args, **kwargs))
+      end
+
+      def upsample(*args, **kwargs)
+        push(MLX::NN::Upsample.new(*args, **kwargs))
+      end
+
+      def method_missing(name, *args, **kwargs, &block)
+        if !@owner.nil? && @owner.respond_to?(name)
+          @owner.public_send(name, *args, **kwargs, &block)
+        else
+          super
+        end
+      end
+
+      def respond_to_missing?(name, include_private = false)
+        (!@owner.nil? && @owner.respond_to?(name, include_private)) || super
+      end
+
+      private
+
+      def collect_modules(&block)
+        previous = @collector
+        @collector = []
+        returned = instance_eval(&block)
+        collected = @collector.dup
+        if collected.empty? && !returned.nil?
+          collected << returned
+        end
+        collected
+      ensure
+        @collector = previous
+      end
+
+      def push(module_obj)
+        @collector << module_obj unless @collector.nil?
+        module_obj
+      end
+
+      def __dsl_modules_from(existing, &block)
+        out = existing.dup
+        out.concat(collect_modules(&block)) if block_given?
+        out.map { |entry| __dsl_normalize_module_entry(entry) }
+      end
+
+      def __dsl_normalize_module_entry(entry)
+        return entry if entry.is_a?(MLX::NN::Module)
+
+        if entry.is_a?(Class)
+          return entry.new if entry <= MLX::NN::Module
+
+          raise TypeError, "builder entries must be MLX::NN::Module instances, MLX::NN::Module classes, or callables"
+        end
+
+        return MLX::DSL::Callable.new(entry) if entry.respond_to?(:call)
+
+        raise TypeError, "builder entries must be MLX::NN::Module instances, MLX::NN::Module classes, or callables"
+      end
+
+      def __dsl_reject_layer_constructor_args!(args, kwargs, entry_type)
+        return if args.empty? && kwargs.empty?
+
+        raise ArgumentError, "layer entry #{entry_type} does not accept constructor arguments"
+      end
+
+      def __dsl_collect_repeated_entries(count, &block)
+        raise ArgumentError, "repeat requires a block" unless block_given?
+
+        repeats = count.to_i
+        raise ArgumentError, "repeat count must be non-negative" if repeats.negative?
+
+        out = []
+        repeats.times do |index|
+          out.concat(
+            collect_modules do
+              __dsl_call_repeat_block(block, index)
+            end
+          )
+        end
+        out
+      end
+
+      def __dsl_call_repeat_block(block, index)
+        return instance_eval(&block) if block.arity.zero?
+
+        block.call(index)
+      end
+
+      def __dsl_build_class_stack_layers(count, layer_class, args, kwargs)
+        repeats = count.to_i
+        raise ArgumentError, "stack count must be non-negative" if repeats.negative?
+        unless layer_class.is_a?(Class) && layer_class <= MLX::NN::Module
+          raise TypeError, "stack layer class must inherit from MLX::NN::Module"
+        end
+
+        Array.new(repeats) { layer_class.new(*args, **kwargs) }
+      end
+    end
+  end
+end