mlx 0.30.7 → 0.30.7.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: ce770b59d71e1d5fbf8697bfebde05006fa52a770f8526575f94944cf161d05b
-  data.tar.gz: 814aa4e6c063b3b36c3b1d3c4c26d8f5fa3952a66947d294a8edb3864b07ece2
+  metadata.gz: 25d582e4816d69b27713a4027534b75cd00ca72557e69681daf07146d3e79ef2
+  data.tar.gz: c010252aa355370a531fa4f3b9bf8cc729876d2f7fb9ae8b8e0d6a1eb6cb57c4
 SHA512:
-  metadata.gz: 3702f6d4445ea4af978ffcebb71dcef91f8e28d4f4d79a1d16a49b66848023a84b7b8ff2a614fa7b2c236d324dfb92b8ccb01e70ff59c53234ac28ee6fc09b39
-  data.tar.gz: 52bf528fee068f422dac611fdca22fae8089a7ea9a3a4d0cd42a21aecc53b5219b0dcf83d79f5fa907a1c99f97acb831d952474963071b566b3483e91c4a4a72
+  metadata.gz: 53e629e845342f173c04c7c6d9d976a29dd5492ae945239897d3168a586288ec958ba58753345317f301220ac5f4b91a22f97731ab799fcea5d59f3d19e48214
+  data.tar.gz: 5b04f2e63e3dcdb6a0282184a310600f4fb72e606b45e8e7a27a7b9461abef3a598afe2eb5525e66457da45b8a84a6c2f07c871c955e9cddd4308971496c7fd1

data/ext/mlx/native.cpp

@@ -6625,7 +6625,8 @@ static VALUE core_clear_cache(VALUE) {
 
 static VALUE core_metal_is_available(VALUE) {
   try {
-    return mxmetal::is_available() ? Qtrue : Qfalse;
+    const mx::Device gpu_device(mx::Device::gpu, 0);
+    return mx::is_available(gpu_device) ? Qtrue : Qfalse;
   } catch (const std::exception& error) {
     raise_std_exception(error);
     return Qnil;
@@ -6654,7 +6655,12 @@ static VALUE core_metal_stop_capture(VALUE) {
 
 static VALUE core_metal_device_info(VALUE) {
   try {
-    const auto& info = mxmetal::device_info();
+    const mx::Device gpu_device(mx::Device::gpu, 0);
+    if (!mx::is_available(gpu_device)) {
+      rb_raise(rb_eRuntimeError, "[metal_device_info] Metal GPU device is not available");
+    }
+
+    const auto& info = mx::device_info(gpu_device);
     VALUE hash = rb_hash_new();
     for (const auto& [key, value] : info) {
       VALUE ruby_key = rb_utf8_str_new(key.c_str(), static_cast<long>(key.size()));
@@ -7884,9 +7890,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 +8026,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/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,
@@ -373,12 +375,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/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

@@ -0,0 +1,385 @@
+# 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 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
+
+      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