mlx 1.0.0

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

@@ -0,0 +1,167 @@
+# frozen_string_literal: true
+
+module MLX
+  module NN
+    class PoolBase < Module
+      def initialize(pooling_symbol, kernel_size, stride, padding, padding_value)
+        super()
+        @pooling_symbol = pooling_symbol
+        @kernel_size = kernel_size
+        @stride = stride
+        @padding = padding
+        @padding_value = padding_value
+      end
+
+      def call(x)
+        if @padding.any? { |p| p[0] > 0 }
+          x = MLX::Core.pad(x, [[0, 0], *@padding, [0, 0]], nil, @padding_value)
+        end
+        windows = sliding_windows(x, @kernel_size, @stride)
+        reduce_windows(windows)
+      end
+
+      private
+
+      def value_or_list(value, n, message)
+        if value.is_a?(Array)
+          unless value.length == n && value.all? { |v| v.is_a?(Integer) }
+            raise ArgumentError, message
+          end
+          return value
+        end
+
+        unless value.is_a?(Integer)
+          raise ArgumentError, message
+        end
+        Array.new(n, value)
+      end
+
+      def reduce_windows(windows)
+        result = windows
+        window_dims = @kernel_size.length
+        window_dims.times do
+          axis = result.ndim - 2
+          result = if @pooling_symbol == :max
+            MLX::Core.max(result, axis)
+          else
+            MLX::Core.mean(result, axis)
+          end
+        end
+        result
+      end
+
+      def sliding_windows(x, window_shape, window_strides)
+        if x.ndim < 3
+          raise ArgumentError,
+                "To extract sliding windows at least 1 spatial dimension (3 total) is needed but the input only has #{x.ndim} dimensions."
+        end
+
+        spatial_dims = x.shape[1...-1]
+        unless spatial_dims.length == window_shape.length && window_shape.length == window_strides.length
+          raise ArgumentError,
+                "To extract sliding windows the window shapes and strides must have the same number of spatial dimensions as the signal."
+        end
+
+        shape = x.shape
+        strides = Array.new(shape.length)
+        running = 1
+        (shape.length - 1).downto(0) do |i|
+          strides[i] = running
+          running *= shape[i]
+        end
+
+        final_shape = [shape[0]]
+        spatial_dims.each_with_index do |size, i|
+          window = window_shape[i]
+          stride = window_strides[i]
+          final_shape << ((size - window) / stride + 1)
+        end
+        final_shape.concat(window_shape)
+        final_shape << shape[-1]
+
+        final_strides = [strides[0]]
+        spatial_dims.each_with_index do |_size, i|
+          final_strides << (strides[i + 1] * window_strides[i])
+        end
+        final_strides.concat(strides[1...-1])
+        final_strides << strides[-1]
+
+        MLX::Core.as_strided(x, final_shape, final_strides)
+      end
+    end
+
+    class Pool1dBase < PoolBase
+      def initialize(pooling_symbol, padding_value, kernel_size, stride = nil, padding = 0)
+        class_name = self.class.name.split("::").last
+        msg = "[#{class_name}] '%s' must be an integer or a tuple containing 1 integer"
+        kernel_size = value_or_list(kernel_size, 1, format(msg, "kernel_size"))
+        stride = stride.nil? ? kernel_size : value_or_list(stride, 1, format(msg, "stride"))
+        padding = value_or_list(padding, 1, format(msg, "padding")).map { |p| [p, p] }
+        super(pooling_symbol, kernel_size, stride, padding, padding_value)
+      end
+    end
+
+    class Pool2dBase < PoolBase
+      def initialize(pooling_symbol, padding_value, kernel_size, stride = nil, padding = 0)
+        class_name = self.class.name.split("::").last
+        msg = "[#{class_name}] '%s' must be an integer or a tuple containing 2 integers"
+        kernel_size = value_or_list(kernel_size, 2, format(msg, "kernel_size"))
+        stride = stride.nil? ? kernel_size : value_or_list(stride, 2, format(msg, "stride"))
+        padding = value_or_list(padding, 2, format(msg, "padding")).map { |p| [p, p] }
+        super(pooling_symbol, kernel_size, stride, padding, padding_value)
+      end
+    end
+
+    class Pool3dBase < PoolBase
+      def initialize(pooling_symbol, padding_value, kernel_size, stride = nil, padding = 0)
+        class_name = self.class.name.split("::").last
+        msg = "[#{class_name}] '%s' must be an integer or a tuple containing 3 integers"
+        kernel_size = value_or_list(kernel_size, 3, format(msg, "kernel_size"))
+        stride = stride.nil? ? kernel_size : value_or_list(stride, 3, format(msg, "stride"))
+        padding = value_or_list(padding, 3, format(msg, "padding")).map { |p| [p, p] }
+        super(pooling_symbol, kernel_size, stride, padding, padding_value)
+      end
+    end
+
+    class MaxPool1d < Pool1dBase
+      def initialize(kernel_size, stride: nil, padding: 0)
+        super(:max, -Float::INFINITY, kernel_size, stride, padding)
+      end
+    end
+
+    class AvgPool1d < Pool1dBase
+      def initialize(kernel_size, stride: nil, padding: 0)
+        super(:mean, 0.0, kernel_size, stride, padding)
+      end
+    end
+
+    class MaxPool2d < Pool2dBase
+      def initialize(kernel_size, stride: nil, padding: 0)
+        super(:max, -Float::INFINITY, kernel_size, stride, padding)
+      end
+    end
+
+    class AvgPool2d < Pool2dBase
+      def initialize(kernel_size, stride: nil, padding: 0)
+        super(:mean, 0.0, kernel_size, stride, padding)
+      end
+    end
+
+    class MaxPool3d < Pool3dBase
+      def initialize(kernel_size, stride: nil, padding: 0)
+        super(:max, -Float::INFINITY, kernel_size, stride, padding)
+      end
+    end
+
+    class AvgPool3d < Pool3dBase
+      def initialize(kernel_size, stride: nil, padding: 0)
+        super(:mean, 0.0, kernel_size, stride, padding)
+      end
+    end
+
+    remove_const(:Pool1dBase)
+    remove_const(:Pool2dBase)
+    remove_const(:Pool3dBase)
+    remove_const(:PoolBase)
+  end
+end

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

@@ -0,0 +1,126 @@
+# frozen_string_literal: true
+
+module MLX
+  module NN
+    class RoPE < Module
+      def initialize(dims, traditional: false, base: 10_000.0, scale: 1.0)
+        super()
+        @dims = dims
+        @traditional = traditional
+        @base = base
+        @scale = scale
+      end
+
+      def call(x, offset: 0)
+        MLX::Core.rope(x, @dims, @traditional, @base, @scale, offset)
+      end
+    end
+
+    class SinusoidalPositionalEncoding < Module
+      def initialize(
+        dims,
+        min_freq: 0.0001,
+        max_freq: 1.0,
+        scale: nil,
+        cos_first: false,
+        full_turns: false
+      )
+        super()
+
+        half_dims = dims / 2
+        one_zero = MLX::Core.subtract(
+          1.0,
+          MLX::Core.divide(
+            MLX::Core.arange(0, half_dims, 1, MLX::Core.float32),
+            (half_dims - 1).to_f
+          )
+        )
+        min_log = Math.log(min_freq)
+        max_log = Math.log(max_freq)
+
+        self._sigmas = MLX::Core.exp(
+          MLX::Core.add(MLX::Core.multiply(one_zero, max_log - min_log), min_log)
+        )
+        self._sigmas = MLX::Core.multiply(_sigmas, 2.0 * Math::PI) if full_turns
+
+        @scale = scale || Math.sqrt(2.0 / dims)
+        @cos_first = cos_first
+      end
+
+      def call(x)
+        y = MLX::Core.multiply(MLX::Core.expand_dims(x, -1), _sigmas)
+        cosy = MLX::Core.cos(y)
+        siny = MLX::Core.sin(y)
+        y = if @cos_first
+          MLX::Core.concatenate([cosy, siny], -1)
+        else
+          MLX::Core.concatenate([siny, cosy], -1)
+        end
+
+        if @scale != 1.0
+          MLX::Core.multiply(y, @scale)
+        else
+          y
+        end
+      end
+    end
+
+    class ALiBi < Module
+      class << self
+        def create_alibi_matrix(
+          q_sequence_length:,
+          k_sequence_length:,
+          num_heads:,
+          offset:,
+          dtype: MLX::Core.float32
+        )
+          x1 = MLX::Core.arange(offset, q_sequence_length, 1)
+          x2 = MLX::Core.arange(0, k_sequence_length, 1)
+          x1_col = MLX::Core.reshape(x1, [x1.shape[0], 1])
+          x2_row = MLX::Core.reshape(x2, [1, x2.shape[0]])
+          distance = MLX::Core.multiply(MLX::Core.abs(MLX::Core.subtract(x1_col, x2_row)), -1.0)
+          distance = MLX::Core.expand_dims(MLX::Core.expand_dims(distance, 0), 1)
+
+          slope = create_alibi_slope(num_heads: num_heads, dtype: dtype)
+          MLX::Core.multiply(distance, slope).astype(dtype)
+        end
+
+        def create_alibi_slope(num_heads:, dtype:)
+          slopes = get_slopes(num_heads)
+          out = MLX::Core.array(slopes, dtype)
+          MLX::Core.expand_dims(MLX::Core.expand_dims(out, -1), -1)
+        end
+
+        private
+
+        def get_slopes(n)
+          if integer_log2?(n)
+            start = 2.0**(-(2.0**(-(Math.log2(n) - 3.0))))
+            Array.new(n) { |i| start * (start**i) }
+          else
+            closest_power_of_2 = 2**Math.log2(n).floor
+            base = get_slopes(closest_power_of_2)
+            extras = get_slopes(2 * closest_power_of_2).each_with_index.select { |_, i| i.even? }.map(&:first)
+            base + extras.first(n - closest_power_of_2)
+          end
+        end
+
+        def integer_log2?(n)
+          Math.log2(n).to_i == Math.log2(n)
+        end
+      end
+
+      def call(attention_scores, offset: 0, mask: nil)
+        alibi_mask = self.class.create_alibi_matrix(
+          q_sequence_length: attention_scores.shape[-2] + offset,
+          k_sequence_length: attention_scores.shape[-1],
+          num_heads: attention_scores.shape[1],
+          offset: offset,
+          dtype: attention_scores.dtype
+        )
+        alibi_mask = MLX::Core.add(alibi_mask, mask) unless mask.nil?
+        MLX::Core.add(attention_scores, alibi_mask)
+      end
+    end
+  end
+end

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

@@ -0,0 +1,215 @@
+# frozen_string_literal: true
+
+module MLX
+  module NN
+    class << self
+      def quantize(
+        model,
+        group_size: nil,
+        bits: nil,
+        mode: "affine",
+        quantize_input: false,
+        class_predicate: nil
+      )
+        class_predicate ||= lambda do |_path, module_obj|
+          module_obj.respond_to?(:to_quantized)
+        end
+
+        maybe_quantize = lambda do |path, module_obj|
+          decision = class_predicate.call(path, module_obj)
+          decision = false if decision.respond_to?(:empty?) && decision.empty?
+
+          return module_obj unless decision
+
+          unless module_obj.respond_to?(:to_quantized)
+            raise ArgumentError, "Unable to quantize model of type #{module_obj.class}"
+          end
+
+          if decision == true
+            kwargs = { group_size: group_size, bits: bits, mode: mode }
+            kwargs[:quantize_input] = true if quantize_input
+            module_obj.to_quantized(**kwargs)
+          elsif decision.is_a?(Hash)
+            kwargs = symbolize_hash(decision.dup)
+            kwargs.delete(:quantize_input) if kwargs[:quantize_input] == false
+            module_obj.to_quantized(**kwargs)
+          else
+            raise ArgumentError, "class_predicate must return a bool or a hash of quantization kwargs"
+          end
+        end
+
+        leaves = model.leaf_modules
+        updated = MLX::Utils.tree_map_with_path(
+          maybe_quantize,
+          leaves,
+          is_leaf: lambda { |v| v.is_a?(MLX::NN::Module) }
+        )
+        model.update_modules(updated)
+        model
+      end
+
+      private
+
+      def defaults_for_mode(mode, group_size, bits)
+        defaults = {
+          "affine" => [64, 4],
+          "mxfp4" => [32, 4],
+          "nvfp4" => [16, 4],
+          "mxfp8" => [32, 8]
+        }
+        default_group_size, default_bits = defaults.fetch(mode.to_s) do
+          raise ArgumentError, "Unsupported quantization mode #{mode}"
+        end
+        [group_size || default_group_size, bits || default_bits]
+      end
+
+      def symbolize_hash(hash)
+        hash.each_with_object({}) do |(key, value), out|
+          out[key.to_sym] = value
+        end
+      end
+    end
+
+    class QuantizedEmbedding < Module
+      attr_reader :group_size, :bits, :mode, :num_embeddings, :dims
+
+      def initialize(num_embeddings, dims, group_size = nil, bits = nil, mode: "affine")
+        super()
+        @group_size, @bits = MLX::NN.__send__(:defaults_for_mode, mode, group_size, bits)
+        @mode = mode
+
+        scale = Math.sqrt(1.0 / dims)
+        weight = MLX::Core.normal([num_embeddings, dims], 0.0, scale)
+        q_weight, q_scales, *q_biases = MLX::Core.quantize(weight, group_size, bits, mode)
+        self.weight = q_weight
+        self.scales = q_scales
+        self.biases = q_biases.empty? ? nil : q_biases[0]
+
+        @num_embeddings = num_embeddings
+        @dims = dims
+        freeze
+      end
+
+      def call(x)
+        gathered_weight = MLX::Core.take(weight, x, 0)
+        gathered_scales = MLX::Core.take(scales, x, 0)
+        gathered_biases = biases.nil? ? nil : MLX::Core.take(biases, x, 0)
+        MLX::Core.dequantize(gathered_weight, gathered_scales, gathered_biases, group_size, bits, mode)
+      end
+
+      def as_linear(x)
+        MLX::Core.quantized_matmul(x, weight, scales, biases, true, group_size, bits, mode)
+      end
+
+      def self.from_embedding(embedding_layer, group_size = nil, bits = nil, mode: "affine")
+        num_embeddings, dims = embedding_layer.weight.shape
+        out = new(num_embeddings, dims, group_size, bits, mode: mode)
+        q_weight, q_scales, *q_biases = MLX::Core.quantize(embedding_layer.weight, group_size, bits, mode)
+        out.weight = q_weight
+        out.scales = q_scales
+        out.biases = q_biases.empty? ? nil : q_biases[0]
+        out
+      end
+    end
+
+    class QuantizedLinear < Module
+      attr_reader :group_size, :bits, :mode
+
+      def initialize(input_dims, output_dims, bias = true, group_size = nil, bits = nil, mode: "affine")
+        super()
+        @group_size, @bits = MLX::NN.__send__(:defaults_for_mode, mode, group_size, bits)
+        @mode = mode
+
+        scale = Math.sqrt(1.0 / input_dims)
+        weight = MLX::Core.uniform([output_dims, input_dims], -scale, scale)
+        q_weight, q_scales, *q_biases = MLX::Core.quantize(weight, group_size, bits, mode)
+        self.weight = q_weight
+        self.scales = q_scales
+        self.biases = q_biases.empty? ? nil : q_biases[0]
+        self.bias = MLX::Core.zeros([output_dims], MLX::Core.float32) if bias
+
+        freeze
+      end
+
+      def call(x)
+        out = MLX::Core.quantized_matmul(x, weight, scales, biases, true, group_size, bits, mode)
+        if state.key?("bias")
+          MLX::Core.add(out, bias)
+        else
+          out
+        end
+      end
+
+      def self.from_linear(linear_layer, group_size = nil, bits = nil, mode: "affine")
+        output_dims, input_dims = linear_layer.weight.shape
+        out = new(input_dims, output_dims, false, group_size, bits, mode: mode)
+        q_weight, q_scales, *q_biases = MLX::Core.quantize(linear_layer.weight, group_size, bits, mode)
+        out.weight = q_weight
+        out.scales = q_scales
+        out.biases = q_biases.empty? ? nil : q_biases[0]
+        out.bias = linear_layer.bias if linear_layer.state.key?("bias")
+        out
+      end
+    end
+
+    class QQLinear < Module
+      attr_reader :group_size, :bits, :mode
+
+      def initialize(input_dims, output_dims, group_size = nil, bits = nil, mode: "nvfp4")
+        super()
+        @group_size, @bits = MLX::NN.__send__(:defaults_for_mode, mode, group_size, bits)
+        @mode = mode
+
+        scale = Math.sqrt(1.0 / input_dims)
+        self.weight = MLX::Core.uniform([output_dims, input_dims], -scale, scale)
+        @_quantized = false
+      end
+
+      def quantize
+        return self if @_quantized
+
+        q_weight, q_scales, *_ = MLX::Core.quantize(weight, group_size, bits, mode)
+        self.weight = q_weight
+        self.scales = q_scales
+        @_quantized = true
+        self
+      end
+
+      def dequantize
+        return self unless @_quantized
+
+        self.weight = MLX::Core.dequantize(weight, scales, nil, group_size, bits, mode)
+        state.delete("scales")
+        @_quantized = false
+        self
+      end
+
+      def train(mode = true)
+        super
+        if training
+          dequantize
+        else
+          quantize
+        end
+        self
+      end
+
+      def call(x)
+        q_scales = state["scales"]
+        MLX::Core.qqmm(x, weight, q_scales, group_size, bits, mode)
+      end
+
+      def self.from_linear(linear_layer, group_size = nil, bits = nil, mode: "nvfp4")
+        if linear_layer.state.key?("bias")
+          raise NotImplementedError, "QQLinear does not support bias yet."
+        end
+
+        output_dims, input_dims = linear_layer.weight.shape
+        out = new(input_dims, output_dims, group_size, bits, mode: mode)
+        out.weight = linear_layer.weight
+        out.train(linear_layer.training)
+        out
+      end
+    end
+  end
+end

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

@@ -0,0 +1,135 @@
+# frozen_string_literal: true
+
+module MLX
+  module NN
+    class RNN < Module
+      def initialize(input_size, hidden_size, bias: true, nonlinearity: nil)
+        super()
+
+        @nonlinearity = nonlinearity || lambda { |z| MLX::NN.tanh(z) }
+        unless @nonlinearity.respond_to?(:call)
+          raise ArgumentError, "Nonlinearity must be callable. Current value: #{nonlinearity}."
+        end
+
+        scale = 1.0 / Math.sqrt(hidden_size)
+        @hidden_size = hidden_size
+        self.Wxh = MLX::Core.uniform([hidden_size, input_size], -scale, scale)
+        self.Whh = MLX::Core.uniform([hidden_size, hidden_size], -scale, scale)
+        self.bias = bias ? MLX::Core.uniform([hidden_size], -scale, scale) : nil
+      end
+
+      def call(x, hidden = nil)
+        x = MLX::Core.matmul(x, self.Wxh.T)
+        x = MLX::Core.add(x, self.bias) unless self.bias.nil?
+
+        all_hidden = []
+        sequence_axis = x.ndim - 2
+        x.shape[sequence_axis].times do |idx|
+          step = MLX::Core.take(x, idx, sequence_axis)
+          step = MLX::Core.add(step, MLX::Core.matmul(hidden, self.Whh.T)) unless hidden.nil?
+          hidden = @nonlinearity.call(step)
+          all_hidden << hidden
+        end
+
+        MLX::Core.stack(all_hidden, -2)
+      end
+    end
+
+    class GRU < Module
+      def initialize(input_size, hidden_size, bias: true)
+        super()
+
+        @hidden_size = hidden_size
+        scale = 1.0 / Math.sqrt(hidden_size)
+        self.Wx = MLX::Core.uniform([3 * hidden_size, input_size], -scale, scale)
+        self.Wh = MLX::Core.uniform([3 * hidden_size, hidden_size], -scale, scale)
+        self.b = bias ? MLX::Core.uniform([3 * hidden_size], -scale, scale) : nil
+        self.bhn = bias ? MLX::Core.uniform([hidden_size], -scale, scale) : nil
+      end
+
+      def call(x, hidden = nil)
+        x = MLX::Core.matmul(x, self.Wx.T)
+        x = MLX::Core.add(x, self.b) unless self.b.nil?
+
+        x_rz, x_n = MLX::Core.split(x, [2 * @hidden_size], x.ndim - 1)
+        all_hidden = []
+        sequence_axis = x.ndim - 2
+
+        x.shape[sequence_axis].times do |idx|
+          rz = MLX::Core.take(x_rz, idx, sequence_axis)
+          h_proj_n = nil
+
+          unless hidden.nil?
+            h_proj = MLX::Core.matmul(hidden, self.Wh.T)
+            h_proj_rz, h_proj_n = MLX::Core.split(h_proj, [2 * @hidden_size], h_proj.ndim - 1)
+            h_proj_n = MLX::Core.add(h_proj_n, self.bhn) unless self.bhn.nil?
+            rz = MLX::Core.add(rz, h_proj_rz)
+          end
+
+          rz = MLX::Core.sigmoid(rz)
+          r, z = MLX::Core.split(rz, 2, rz.ndim - 1)
+          n = MLX::Core.take(x_n, idx, sequence_axis)
+          n = MLX::Core.add(n, MLX::Core.multiply(r, h_proj_n)) unless hidden.nil?
+          n = MLX::Core.tanh(n)
+
+          if hidden.nil?
+            hidden = MLX::Core.multiply(MLX::Core.subtract(1.0, z), n)
+          else
+            hidden = MLX::Core.add(
+              MLX::Core.multiply(MLX::Core.subtract(1.0, z), n),
+              MLX::Core.multiply(z, hidden)
+            )
+          end
+
+          all_hidden << hidden
+        end
+
+        MLX::Core.stack(all_hidden, -2)
+      end
+    end
+
+    class LSTM < Module
+      def initialize(input_size, hidden_size, bias: true)
+        super()
+
+        @hidden_size = hidden_size
+        scale = 1.0 / Math.sqrt(hidden_size)
+        self.Wx = MLX::Core.uniform([4 * hidden_size, input_size], -scale, scale)
+        self.Wh = MLX::Core.uniform([4 * hidden_size, hidden_size], -scale, scale)
+        self.bias = bias ? MLX::Core.uniform([4 * hidden_size], -scale, scale) : nil
+      end
+
+      def call(x, hidden = nil, cell = nil)
+        x = MLX::Core.matmul(x, self.Wx.T)
+        x = MLX::Core.add(x, self.bias) unless self.bias.nil?
+
+        all_hidden = []
+        all_cell = []
+        sequence_axis = x.ndim - 2
+
+        x.shape[sequence_axis].times do |idx|
+          ifgo = MLX::Core.take(x, idx, sequence_axis)
+          ifgo = MLX::Core.add(ifgo, MLX::Core.matmul(hidden, self.Wh.T)) unless hidden.nil?
+
+          i, f, g, o = MLX::Core.split(ifgo, 4, ifgo.ndim - 1)
+          i = MLX::Core.sigmoid(i)
+          f = MLX::Core.sigmoid(f)
+          g = MLX::Core.tanh(g)
+          o = MLX::Core.sigmoid(o)
+
+          cell = if cell.nil?
+            MLX::Core.multiply(i, g)
+          else
+            MLX::Core.add(MLX::Core.multiply(f, cell), MLX::Core.multiply(i, g))
+          end
+          hidden = MLX::Core.multiply(o, MLX::Core.tanh(cell))
+
+          all_cell << cell
+          all_hidden << hidden
+        end
+
+        [MLX::Core.stack(all_hidden, -2), MLX::Core.stack(all_cell, -2)]
+      end
+    end
+  end
+end