mlx 0.30.7

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

@@ -0,0 +1,309 @@
+# frozen_string_literal: true
+
+module MLX
+  module NN
+    class << self
+      def sum_gradients(group = nil)
+        group ||= MLX::Core.init(false, "any")
+        if group.size == 1
+          ->(x) { x }
+        else
+          ->(x) { MLX::Core.all_sum(x, group) }
+        end
+      end
+
+      def shard_inplace(module_obj, sharding, segments: 1, group: nil)
+        group ||= MLX::Core.init(false, "any")
+
+        sharding_fn = if sharding.is_a?(String)
+          check_sharding(sharding)
+          sharding == "all-to-sharded" ? all_to_sharded(segments) : sharded_to_all(segments)
+        else
+          sharding
+        end
+
+        module_obj.update(shard(module_obj.parameters, sharding_fn, group))
+        module_obj
+      end
+
+      def shard_linear(module_obj, sharding, segments: 1, group: nil)
+        check_sharding(sharding)
+        fns = {
+          ["all-to-sharded", true] => ->(m, s, g) { AllToShardedLinear.from_linear(m, segments: s, group: g) },
+          ["all-to-sharded", false] => ->(m, s, g) { QuantizedAllToShardedLinear.from_quantized_linear(m, segments: s, group: g) },
+          ["sharded-to-all", true] => ->(m, s, g) { ShardedToAllLinear.from_linear(m, segments: s, group: g) },
+          ["sharded-to-all", false] => lambda { |m, s, g|
+            QuantizedShardedToAllLinear.from_quantized_linear(m, segments: s, group: g)
+          }
+        }
+        fns.fetch([sharding, module_obj.is_a?(Linear)]).call(module_obj, segments, group)
+      end
+
+      private
+
+      def split_segments(weight, segments, axis)
+        if segments.is_a?(Integer) || (segments.is_a?(Array) && !segments.empty? && segments[0].is_a?(Integer))
+          return MLX::Core.split(weight, segments, axis)
+        end
+
+        n = weight.shape[axis]
+        indices = segments.map { |s| (s * n).to_i }
+        MLX::Core.split(weight, indices, axis)
+      end
+
+      def shard(parameters, sharding_predicate, group)
+        n = group.size
+        r = group.rank
+        return parameters if n == 1
+
+        shard_fn = lambda do |path, weight|
+          return weight unless weight.is_a?(MLX::Core::Array)
+
+          spec = sharding_predicate.call(path.to_s, weight)
+          return weight if spec.nil?
+
+          axis = nil
+          segments = 1
+          if spec.is_a?(Integer)
+            axis = spec
+          elsif spec.is_a?(Array) && spec.length == 2
+            axis = spec[0]
+            segments = spec[1]
+          else
+            raise ArgumentError, "The sharding function should return int or tuple [axis, segments]"
+          end
+
+          part_segments = split_segments(weight, segments, axis)
+          shard_parts = part_segments.map do |part|
+            split_segments(part, n, axis)[r]
+          end
+          MLX::Core.concatenate(shard_parts, axis)
+        end
+
+        MLX::Utils.tree_map_with_path(shard_fn, parameters)
+      end
+
+      def all_to_sharded(segments)
+        lambda do |path, weight|
+          if path.end_with?("bias")
+            [-1, segments]
+          else
+            [[weight.ndim - 2, 0].max, segments]
+          end
+        end
+      end
+
+      def sharded_to_all(segments)
+        lambda do |path, _weight|
+          if path.end_with?("bias")
+            nil
+          else
+            [-1, segments]
+          end
+        end
+      end
+
+      def check_sharding(sharding)
+        return if %w[all-to-sharded sharded-to-all].include?(sharding)
+
+        raise ArgumentError,
+              "Sharding type sharding=#{sharding} not supported, choose one of 'all-to-sharded' or 'sharded-to-all'"
+      end
+    end
+
+    class AllToShardedLinear < Module
+      def initialize(input_dims, output_dims, bias: true, group: nil)
+        super()
+        @group = group || MLX::Core.init(false, "any")
+        n = @group.size
+        if (output_dims % n) != 0
+          raise ArgumentError, "Cannot shard the output of size #{output_dims} across #{n} devices."
+        end
+
+        scale = Math.sqrt(1.0 / input_dims)
+        self.weight = MLX::Core.uniform([output_dims / n, input_dims], -scale, scale)
+        self.bias = MLX::Core.uniform([output_dims / n], -scale, scale) if bias
+      end
+
+      def call(x)
+        x = MLX::NN.sum_gradients(@group).call(x)
+        out = MLX::Core.matmul(x, weight.T)
+        state.key?("bias") ? MLX::Core.add(out, bias) : out
+      end
+
+      def self.from_linear(linear_layer, segments: 1, group: nil)
+        group ||= MLX::Core.init(false, "any")
+        output_dims, input_dims = linear_layer.weight.shape
+        sl = new(input_dims, output_dims, bias: linear_layer.state.key?("bias"), group: group)
+        sl.update(
+          MLX::NN.__send__(
+            :shard,
+            linear_layer.parameters,
+            MLX::NN.__send__(:all_to_sharded, segments),
+            group
+          )
+        )
+        sl
+      end
+    end
+
+    class ShardedToAllLinear < Module
+      def initialize(input_dims, output_dims, bias: true, group: nil)
+        super()
+        @group = group || MLX::Core.init(false, "any")
+        n = @group.size
+        if (input_dims % n) != 0
+          raise ArgumentError, "The input of size #{input_dims} cannot be sharded across #{n} devices."
+        end
+
+        scale = Math.sqrt(1.0 / input_dims)
+        self.weight = MLX::Core.uniform([output_dims, input_dims / n], -scale, scale)
+        self.bias = MLX::Core.uniform([output_dims], -scale, scale) if bias
+      end
+
+      def call(x)
+        out = MLX::Core.matmul(x, weight.T)
+        out = MLX::Core.all_sum(out, @group)
+        state.key?("bias") ? MLX::Core.add(out, bias) : out
+      end
+
+      def self.from_linear(linear_layer, segments: 1, group: nil)
+        group ||= MLX::Core.init(false, "any")
+        output_dims, input_dims = linear_layer.weight.shape
+        sl = new(input_dims, output_dims, bias: linear_layer.state.key?("bias"), group: group)
+        sl.update(
+          MLX::NN.__send__(
+            :shard,
+            linear_layer.parameters,
+            MLX::NN.__send__(:sharded_to_all, segments),
+            group
+          )
+        )
+        sl
+      end
+    end
+
+    class QuantizedAllToShardedLinear < Module
+      attr_reader :group_size, :bits
+
+      def initialize(input_dims, output_dims, bias: true, group_size: 64, bits: 4, group: nil)
+        super()
+        @group = group || MLX::Core.init(false, "any")
+        @group_size = group_size
+        @bits = bits
+
+        n = @group.size
+        if (output_dims % n) != 0
+          raise ArgumentError, "Cannot shard the output of size #{output_dims} across #{n} devices."
+        end
+
+        scale = Math.sqrt(1.0 / input_dims)
+        weight = MLX::Core.uniform([output_dims / n, input_dims], -scale, scale)
+        q_weight, q_scales, *q_biases = MLX::Core.quantize(weight, group_size, bits)
+        self.weight = q_weight
+        self.scales = q_scales
+        self.biases = q_biases.empty? ? nil : q_biases[0]
+        self.bias = MLX::Core.zeros([output_dims / n], MLX::Core.float32) if bias
+
+        freeze
+      end
+
+      def unfreeze(*args, **kwargs)
+        super(*args, **kwargs)
+        freeze(recurse: false)
+      end
+
+      def call(x)
+        x = MLX::NN.sum_gradients(@group).call(x)
+        out = MLX::Core.quantized_matmul(x, weight, scales, self.biases, true, @group_size, @bits)
+        state.key?("bias") ? MLX::Core.add(out, bias) : out
+      end
+
+      def self.from_quantized_linear(quantized_linear_layer, segments: 1, group: nil)
+        group ||= MLX::Core.init(false, "any")
+        output_dims, input_packed = quantized_linear_layer.weight.shape
+        input_dims = (input_packed * 32) / quantized_linear_layer.bits
+
+        sl = new(
+          input_dims,
+          output_dims,
+          bias: quantized_linear_layer.state.key?("bias"),
+          group_size: quantized_linear_layer.group_size,
+          bits: quantized_linear_layer.bits,
+          group: group
+        )
+        sl.update(
+          MLX::NN.__send__(
+            :shard,
+            quantized_linear_layer.parameters,
+            MLX::NN.__send__(:all_to_sharded, segments),
+            group
+          )
+        )
+        sl
+      end
+    end
+
+    class QuantizedShardedToAllLinear < Module
+      attr_reader :group_size, :bits
+
+      def initialize(input_dims, output_dims, bias: true, group_size: 64, bits: 4, group: nil)
+        super()
+        @group = group || MLX::Core.init(false, "any")
+        @group_size = group_size
+        @bits = bits
+
+        n = @group.size
+        if (input_dims % n) != 0
+          raise ArgumentError, "The input of size #{input_dims} cannot be sharded across #{n} devices."
+        end
+
+        scale = Math.sqrt(1.0 / input_dims)
+        weight = MLX::Core.uniform([output_dims, input_dims / n], -scale, scale)
+        q_weight, q_scales, *q_biases = MLX::Core.quantize(weight, group_size, bits)
+        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 unfreeze(*args, **kwargs)
+        super(*args, **kwargs)
+        freeze(recurse: false)
+      end
+
+      def call(x)
+        out = MLX::Core.quantized_matmul(x, weight, scales, self.biases, true, @group_size, @bits)
+        out = MLX::Core.all_sum(out, @group)
+        state.key?("bias") ? MLX::Core.add(out, bias) : out
+      end
+
+      def self.from_quantized_linear(quantized_linear_layer, segments: 1, group: nil)
+        group ||= MLX::Core.init(false, "any")
+        output_dims, input_packed = quantized_linear_layer.weight.shape
+        input_dims = (input_packed * 32) / quantized_linear_layer.bits
+
+        sl = new(
+          input_dims,
+          output_dims,
+          bias: quantized_linear_layer.state.key?("bias"),
+          group_size: quantized_linear_layer.group_size,
+          bits: quantized_linear_layer.bits,
+          group: group
+        )
+        sl.update(
+          MLX::NN.__send__(
+            :shard,
+            quantized_linear_layer.parameters,
+            MLX::NN.__send__(:sharded_to_all, segments),
+            group
+          )
+        )
+        sl
+      end
+    end
+
+  end
+end

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

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+module MLX
+  module NN
+    class Dropout < Module
+      def initialize(p = 0.5)
+        super()
+        unless p >= 0.0 && p < 1.0
+          raise ArgumentError, "The dropout probability #{p} is not in [0, 1)"
+        end
+
+        @p_keep = 1.0 - p
+      end
+
+      def call(x)
+        return x if @p_keep == 1.0 || !training
+
+        mask = MLX::Core.bernoulli(@p_keep, x.shape)
+        MLX::Core.multiply(MLX::Core.multiply(mask, x), 1.0 / @p_keep)
+      end
+    end
+
+    class Dropout2d < Module
+      def initialize(p = 0.5)
+        super()
+        unless p >= 0.0 && p < 1.0
+          raise ArgumentError, "The dropout probability #{p} is not in [0, 1)"
+        end
+
+        @p_keep = 1.0 - p
+      end
+
+      def call(x)
+        unless [3, 4].include?(x.ndim)
+          raise ArgumentError, "Received input with #{x.ndim} dimensions. Expected 3 or 4 dimensions."
+        end
+
+        return x if @p_keep == 1.0 || !training
+
+        mask_shape = x.shape.dup
+        mask_shape[-2] = 1
+        mask_shape[-3] = 1
+        mask = MLX::Core.bernoulli(@p_keep, mask_shape)
+        MLX::Core.multiply(MLX::Core.multiply(mask, x), 1.0 / @p_keep)
+      end
+    end
+
+    class Dropout3d < Module
+      def initialize(p = 0.5)
+        super()
+        unless p >= 0.0 && p < 1.0
+          raise ArgumentError, "The dropout probability #{p} is not in [0, 1)"
+        end
+
+        @p_keep = 1.0 - p
+      end
+
+      def call(x)
+        unless [4, 5].include?(x.ndim)
+          raise ArgumentError, "Received input with #{x.ndim} dimensions. Expected 4 or 5 dimensions."
+        end
+
+        return x if @p_keep == 1.0 || !training
+
+        mask_shape = x.shape.dup
+        mask_shape[-2] = 1
+        mask_shape[-3] = 1
+        mask_shape[-4] = 1
+        mask = MLX::Core.bernoulli(@p_keep, mask_shape)
+        MLX::Core.multiply(MLX::Core.multiply(mask, x), 1.0 / @p_keep)
+      end
+    end
+
+  end
+end

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

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module MLX
+  module NN
+    class Embedding < Module
+      def initialize(num_embeddings, dims)
+        super()
+        scale = Math.sqrt(1.0 / dims)
+        self.weight = MLX::Core.normal([num_embeddings, dims], 0.0, scale)
+      end
+
+      def call(x)
+        MLX::Core.take(weight, x, 0)
+      end
+
+      def as_linear(x)
+        MLX::Core.matmul(x, weight.T)
+      end
+
+      def to_quantized(group_size: nil, bits: nil, mode: "affine", quantize_input: false)
+        raise ArgumentError, "Quantized input is not supported." if quantize_input
+
+        QuantizedEmbedding.from_embedding(self, group_size, bits, mode: mode)
+      end
+    end
+
+  end
+end

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

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+module MLX
+  module NN
+    class Identity < Module
+      def initialize(*args, **kwargs)
+        super()
+        _ = [args, kwargs]
+      end
+
+      def call(x, *args, **kwargs)
+        _ = [args, kwargs]
+        x
+      end
+    end
+
+    class Linear < Module
+      def initialize(input_dims, output_dims, bias: true)
+        super()
+        scale = Math.sqrt(1.0 / input_dims)
+        self.weight = MLX::Core.uniform([output_dims, input_dims], -scale, scale)
+        self.bias = MLX::Core.uniform([output_dims], -scale, scale) if bias
+      end
+
+      def call(x)
+        out = MLX::Core.matmul(x, weight.T)
+        if state.key?("bias")
+          MLX::Core.add(out, bias)
+        else
+          out
+        end
+      end
+
+      def to_quantized(group_size: nil, bits: nil, mode: "affine", quantize_input: false)
+        if quantize_input
+          unless %w[nvfp4 mxfp8].include?(mode.to_s)
+            raise ArgumentError,
+                  "Quantized activations are only supported for 'nvfp4' and 'mxfp8' modes, got #{mode}."
+          end
+
+          QQLinear.from_linear(self, group_size, bits, mode: mode)
+        else
+          QuantizedLinear.from_linear(self, group_size, bits, mode: mode)
+        end
+      end
+    end
+
+    class Bilinear < Module
+      def initialize(input1_dims, input2_dims, output_dims, bias: true)
+        super()
+        scale = Math.sqrt(1.0 / input1_dims)
+        self.weight = MLX::Core.uniform([output_dims, input2_dims, input1_dims], -scale, scale)
+        self.bias = MLX::Core.uniform([output_dims], -scale, scale) if bias
+      end
+
+      def call(x1, x2)
+        out_dims, in2_dims, in1_dims = weight.shape
+
+        x_shape = x1.shape[0...-1]
+        batch = x1.size / in1_dims
+        x1_2d = MLX::Core.reshape(x1, [batch, in1_dims])
+        x2_3d = MLX::Core.reshape(x2, [batch, 1, in2_dims])
+
+        w = MLX::Core.reshape(weight, [out_dims * in2_dims, in1_dims])
+        y = MLX::Core.matmul(x1_2d, w.T)
+        y = MLX::Core.reshape(y, [batch, out_dims, in2_dims])
+        y = MLX::Core.swapaxes(y, -2, -1)
+        y = MLX::Core.matmul(x2_3d, y)
+        y = MLX::Core.squeeze(y, 1)
+
+        out_shape = x_shape.empty? ? [out_dims] : x_shape + [out_dims]
+        y = MLX::Core.reshape(y, out_shape)
+        y = MLX::Core.add(y, bias) if state.key?("bias")
+        y
+      end
+    end
+
+  end
+end

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

@@ -0,0 +1,216 @@
+# frozen_string_literal: true
+
+module MLX
+  module NN
+    class InstanceNorm < Module
+      def initialize(dims, eps: 1e-5, affine: false)
+        super()
+        if affine
+          self.weight = MLX::Core.ones([dims], MLX::Core.float32)
+          self.bias = MLX::Core.zeros([dims], MLX::Core.float32)
+        end
+        @dims = dims
+        @eps = eps
+      end
+
+      def call(x)
+        reduction_axes = (1...(x.ndim - 1)).to_a
+        mean = reduce_mean_axes(x, reduction_axes, keepdims: true)
+        var = MLX::Core.var(x, reduction_axes, true)
+        out = MLX::Core.multiply(
+          MLX::Core.subtract(x, mean),
+          MLX::Core.rsqrt(MLX::Core.add(var, @eps))
+        )
+        if state.key?("weight")
+          MLX::Core.add(MLX::Core.multiply(weight, out), bias)
+        else
+          out
+        end
+      end
+
+      private
+
+      def reduce_mean_axes(x, axes, keepdims: false)
+        return x if axes.empty?
+
+        result = x
+        axes.sort.each_with_index do |axis, i|
+          result = MLX::Core.mean(result, axis - i)
+        end
+        if keepdims
+          axes.sort.each do |axis|
+            result = MLX::Core.expand_dims(result, axis)
+          end
+        end
+        result
+      end
+    end
+
+    class LayerNorm < Module
+      def initialize(dims, eps: 1e-5, affine: true, bias: true)
+        super()
+        if affine
+          self.weight = MLX::Core.ones([dims], MLX::Core.float32)
+          self.bias = MLX::Core.zeros([dims], MLX::Core.float32) if bias
+        end
+        @eps = eps
+        @dims = dims
+      end
+
+      def call(x)
+        w = state.key?("weight") ? weight : nil
+        b = state.key?("bias") ? bias : nil
+        MLX::Core.layer_norm(x, w, b, @eps)
+      end
+    end
+
+    class RMSNorm < Module
+      def initialize(dims, eps: 1e-5)
+        super()
+        self.weight = MLX::Core.ones([dims], MLX::Core.float32)
+        @eps = eps
+      end
+
+      def call(x)
+        MLX::Core.rms_norm(x, weight, @eps)
+      end
+    end
+
+    class GroupNorm < Module
+      def initialize(num_groups, dims, eps: 1e-5, affine: true, pytorch_compatible: false)
+        super()
+        if affine
+          self.bias = MLX::Core.zeros([dims], MLX::Core.float32)
+          self.weight = MLX::Core.ones([dims], MLX::Core.float32)
+        end
+        @num_groups = num_groups
+        @dims = dims
+        @eps = eps
+        @pytorch_compatible = pytorch_compatible
+      end
+
+      def call(x)
+        out = if @pytorch_compatible
+          pytorch_compatible_group_norm(x)
+        else
+          group_norm(x)
+        end
+        if state.key?("weight")
+          MLX::Core.add(MLX::Core.multiply(weight, out), bias)
+        else
+          out
+        end
+      end
+
+      private
+
+      def pytorch_compatible_group_norm(x)
+        batch = x.shape[0]
+        rest = x.shape[1...-1]
+        dims = x.shape[-1]
+        group_size = dims / @num_groups
+        feature_count = rest.reduce(1) { |acc, v| acc * v }
+
+        out = MLX::Core.reshape(x, [batch, feature_count, @num_groups, group_size])
+        out = MLX::Core.transpose(out, [0, 2, 1, 3])
+        out = MLX::Core.reshape(out, [batch, @num_groups, feature_count * group_size])
+        out = MLX::Core.layer_norm(out, nil, nil, @eps)
+        out = MLX::Core.reshape(out, [batch, @num_groups, feature_count, group_size])
+        out = MLX::Core.transpose(out, [0, 2, 1, 3])
+        MLX::Core.reshape(out, [batch, *rest, dims])
+      end
+
+      def group_norm(x)
+        batch = x.shape[0]
+        rest = x.shape[1...-1]
+        dims = x.shape[-1]
+        grouped = x.size / (batch * @num_groups)
+
+        out = MLX::Core.reshape(x, [batch, grouped, @num_groups])
+        means = MLX::Core.expand_dims(MLX::Core.mean(out, 1), 1)
+        var = MLX::Core.var(out, 1, true)
+        out = MLX::Core.multiply(
+          MLX::Core.subtract(out, means),
+          MLX::Core.rsqrt(MLX::Core.add(var, @eps))
+        )
+        MLX::Core.reshape(out, [batch, *rest, dims])
+      end
+    end
+
+    class BatchNorm < Module
+      def initialize(num_features, eps: 1e-5, momentum: 0.1, affine: true, track_running_stats: true)
+        super()
+
+        @num_features = num_features
+        @eps = eps
+        @momentum = momentum
+        @track_running_stats = track_running_stats
+
+        if affine
+          self.weight = MLX::Core.ones([num_features], MLX::Core.float32)
+          self.bias = MLX::Core.zeros([num_features], MLX::Core.float32)
+        end
+
+        if @track_running_stats
+          self.running_mean = MLX::Core.zeros([num_features], MLX::Core.float32)
+          self.running_var = MLX::Core.ones([num_features], MLX::Core.float32)
+          freeze(keys: %w[running_mean running_var], recurse: false)
+        end
+      end
+
+      def unfreeze(*args, **kwargs)
+        super(*args, **kwargs)
+        freeze(keys: %w[running_mean running_var], recurse: false) if @track_running_stats
+      end
+
+      def call(x)
+        if x.ndim < 2 || x.ndim > 4
+          raise ArgumentError, "Expected input tensor to have 2, 3 or 4 dimensions, but got #{x.ndim}"
+        end
+
+        mean, var = calc_stats(x)
+        if training && @track_running_stats
+          mu = @momentum
+          self.running_mean = MLX::Core.add(
+            MLX::Core.multiply(1.0 - mu, running_mean),
+            MLX::Core.multiply(mu, mean)
+          )
+          self.running_var = MLX::Core.add(
+            MLX::Core.multiply(1.0 - mu, running_var),
+            MLX::Core.multiply(mu, var)
+          )
+        elsif @track_running_stats
+          mean = running_mean
+          var = running_var
+        end
+
+        out = MLX::Core.multiply(
+          MLX::Core.subtract(x, mean),
+          MLX::Core.rsqrt(MLX::Core.add(var, @eps))
+        )
+        if state.key?("weight")
+          MLX::Core.add(MLX::Core.multiply(weight, out), bias)
+        else
+          out
+        end
+      end
+
+      private
+
+      def calc_stats(x)
+        reduction_axes = (0...(x.ndim - 1)).to_a
+        mean = reduce_mean_axes(x, reduction_axes)
+        var = MLX::Core.var(x, reduction_axes)
+        [mean, var]
+      end
+
+      def reduce_mean_axes(x, axes)
+        result = x
+        axes.sort.each_with_index do |axis, i|
+          result = MLX::Core.mean(result, axis - i)
+        end
+        result
+      end
+    end
+  end
+end