mlx 1.0.0

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

@@ -0,0 +1,330 @@
+# frozen_string_literal: true
+
+module MLX
+  module NN
+    class MultiHeadAttention < Module
+      attr_reader :num_heads
+
+      def initialize(
+        dims,
+        num_heads,
+        query_input_dims: nil,
+        key_input_dims: nil,
+        value_input_dims: nil,
+        value_dims: nil,
+        value_output_dims: nil,
+        bias: false
+      )
+        super()
+
+        if (dims % num_heads) != 0
+          raise ArgumentError,
+                "The input feature dimensions should be divisible by the number of heads (#{dims} % #{num_heads}) != 0"
+        end
+
+        query_input_dims ||= dims
+        key_input_dims ||= dims
+        value_input_dims ||= key_input_dims
+        value_dims ||= dims
+        value_output_dims ||= dims
+
+        @num_heads = num_heads
+        self.query_proj = Linear.new(query_input_dims, dims, bias: bias)
+        self.key_proj = Linear.new(key_input_dims, dims, bias: bias)
+        self.value_proj = Linear.new(value_input_dims, value_dims, bias: bias)
+        self.out_proj = Linear.new(value_dims, value_output_dims, bias: bias)
+      end
+
+      def call(queries, keys, values, mask = nil)
+        queries, q_was_2d = maybe_batch(queries)
+        keys, = maybe_batch(keys)
+        values, = maybe_batch(values)
+
+        queries = query_proj.call(queries)
+        keys = key_proj.call(keys)
+        values = value_proj.call(values)
+
+        queries = split_heads(queries)
+        keys = split_heads(keys)
+        values = split_heads(values)
+
+        scale = Math.sqrt(1.0 / queries.shape[-1])
+        output = MLX::Core.scaled_dot_product_attention(queries, keys, values, scale, mask)
+        output = MLX::Core.transpose(output, [0, 2, 1, 3])
+        output = output.flatten(-2, -1)
+        output = out_proj.call(output)
+        q_was_2d ? MLX::Core.squeeze(output, 0) : output
+      end
+
+      def self.create_additive_causal_mask(n, dtype = MLX::Core.float32)
+        indices = MLX::Core.arange(0, n, 1)
+        lhs = MLX::Core.reshape(indices, [n, 1])
+        rhs = MLX::Core.reshape(indices, [1, n])
+        mask = MLX::Core.less(lhs, rhs).astype(dtype)
+        MLX::Core.multiply(mask, MLX::Core.finfo(dtype).min)
+      end
+
+      private
+
+      def split_heads(x)
+        batch, length, dims = x.shape
+        head_dim = dims / @num_heads
+        x = MLX::Core.reshape(x, [batch, length, @num_heads, head_dim])
+        MLX::Core.transpose(x, [0, 2, 1, 3])
+      end
+
+      def maybe_batch(x)
+        if x.ndim == 2
+          [MLX::Core.expand_dims(x, 0), true]
+        else
+          [x, false]
+        end
+      end
+    end
+
+    class TransformerEncoderLayer < Module
+      def initialize(
+        dims,
+        num_heads,
+        mlp_dims: nil,
+        dropout: 0.0,
+        activation: nil,
+        norm_first: true
+      )
+        super()
+        mlp_dims ||= dims * 4
+        activation ||= lambda { |x| MLX::NN.relu(x) }
+
+        self.attention = MultiHeadAttention.new(dims, num_heads)
+        self.ln1 = LayerNorm.new(dims)
+        self.ln2 = LayerNorm.new(dims)
+        self.linear1 = Linear.new(dims, mlp_dims)
+        self.linear2 = Linear.new(mlp_dims, dims)
+        self.dropout1 = Dropout.new(dropout)
+        self.dropout2 = Dropout.new(dropout)
+        @activation = activation
+        @norm_first = norm_first
+      end
+
+      def call(x, mask)
+        if @norm_first
+          y = ln1.call(x)
+          y = attention.call(y, y, y, mask)
+          y = dropout1.call(y)
+          x = MLX::Core.add(x, y)
+
+          y = ln2.call(x)
+          y = linear1.call(y)
+          y = @activation.call(y)
+          y = dropout2.call(y)
+          y = linear2.call(y)
+          y = MLX::Core.add(x, y)
+        else
+          y = attention.call(x, x, x, mask)
+          y = dropout1.call(y)
+          x = ln1.call(MLX::Core.add(x, y))
+
+          y = linear1.call(x)
+          y = @activation.call(y)
+          y = dropout2.call(y)
+          y = linear2.call(y)
+          y = ln2.call(MLX::Core.add(x, y))
+        end
+
+        y
+      end
+    end
+
+    class TransformerEncoder < Module
+      def initialize(
+        num_layers,
+        dims,
+        num_heads,
+        mlp_dims: nil,
+        dropout: 0.0,
+        activation: nil,
+        norm_first: true,
+        checkpoint: false
+      )
+        super()
+        activation ||= lambda { |x| MLX::NN.relu(x) }
+        self.layers = Array.new(num_layers) do
+          TransformerEncoderLayer.new(
+            dims,
+            num_heads,
+            mlp_dims: mlp_dims,
+            dropout: dropout,
+            activation: activation,
+            norm_first: norm_first
+          )
+        end
+        self.ln = LayerNorm.new(dims)
+        @checkpoint = checkpoint
+        @layer_fns = if @checkpoint
+          layers.map { |layer| MLX::NN.checkpoint(layer, ->(a, b) { layer.call(a, b) }) }
+        else
+          layers
+        end
+      end
+
+      def call(x, mask)
+        @layer_fns.each do |layer_fn|
+          x = layer_fn.call(x, mask)
+        end
+        ln.call(x)
+      end
+    end
+
+    class TransformerDecoderLayer < Module
+      def initialize(
+        dims,
+        num_heads,
+        mlp_dims: nil,
+        dropout: 0.0,
+        activation: nil,
+        norm_first: true
+      )
+        super()
+        mlp_dims ||= dims * 4
+        activation ||= lambda { |x| MLX::NN.relu(x) }
+
+        self.self_attention = MultiHeadAttention.new(dims, num_heads)
+        self.cross_attention = MultiHeadAttention.new(dims, num_heads)
+        self.ln1 = LayerNorm.new(dims)
+        self.ln2 = LayerNorm.new(dims)
+        self.ln3 = LayerNorm.new(dims)
+        self.linear1 = Linear.new(dims, mlp_dims)
+        self.linear2 = Linear.new(mlp_dims, dims)
+        self.dropout1 = Dropout.new(dropout)
+        self.dropout2 = Dropout.new(dropout)
+        self.dropout3 = Dropout.new(dropout)
+        @activation = activation
+        @norm_first = norm_first
+      end
+
+      def call(x, memory, x_mask, memory_mask)
+        if @norm_first
+          y = ln1.call(x)
+          y = self_attention.call(y, y, y, x_mask)
+          y = dropout1.call(y)
+          x = MLX::Core.add(x, y)
+
+          y = ln2.call(x)
+          y = cross_attention.call(y, memory, memory, memory_mask)
+          y = dropout2.call(y)
+          x = MLX::Core.add(x, y)
+
+          y = ln3.call(x)
+          y = linear1.call(y)
+          y = @activation.call(y)
+          y = dropout3.call(y)
+          y = linear2.call(y)
+          y = MLX::Core.add(x, y)
+        else
+          y = self_attention.call(x, x, x, x_mask)
+          y = dropout1.call(y)
+          x = ln1.call(MLX::Core.add(x, y))
+
+          y = cross_attention.call(y, memory, memory, memory_mask)
+          y = dropout2.call(y)
+          x = ln2.call(MLX::Core.add(x, y))
+
+          y = linear1.call(x)
+          y = @activation.call(y)
+          y = dropout3.call(y)
+          y = linear2.call(y)
+          y = ln3.call(MLX::Core.add(x, y))
+        end
+
+        y
+      end
+    end
+
+    class TransformerDecoder < Module
+      def initialize(
+        num_layers,
+        dims,
+        num_heads,
+        mlp_dims: nil,
+        dropout: 0.0,
+        activation: nil,
+        norm_first: true,
+        checkpoint: false
+      )
+        super()
+        activation ||= lambda { |x| MLX::NN.relu(x) }
+        self.layers = Array.new(num_layers) do
+          TransformerDecoderLayer.new(
+            dims,
+            num_heads,
+            mlp_dims: mlp_dims,
+            dropout: dropout,
+            activation: activation,
+            norm_first: norm_first
+          )
+        end
+        self.ln = LayerNorm.new(dims)
+        @checkpoint = checkpoint
+        @layer_fns = if @checkpoint
+          layers.map do |layer|
+            MLX::NN.checkpoint(layer, ->(a, b, c, d) { layer.call(a, b, c, d) })
+          end
+        else
+          layers
+        end
+      end
+
+      def call(x, memory, x_mask, memory_mask)
+        @layer_fns.each do |layer_fn|
+          x = layer_fn.call(x, memory, x_mask, memory_mask)
+        end
+        ln.call(x)
+      end
+    end
+
+    class Transformer < Module
+      def initialize(
+        dims: 512,
+        num_heads: 8,
+        num_encoder_layers: 6,
+        num_decoder_layers: 6,
+        mlp_dims: nil,
+        dropout: 0.0,
+        activation: nil,
+        custom_encoder: nil,
+        custom_decoder: nil,
+        norm_first: true,
+        checkpoint: false
+      )
+        super()
+
+        activation ||= lambda { |x| MLX::NN.relu(x) }
+        self.encoder = custom_encoder || TransformerEncoder.new(
+          num_encoder_layers,
+          dims,
+          num_heads,
+          mlp_dims: mlp_dims,
+          dropout: dropout,
+          activation: activation,
+          norm_first: norm_first,
+          checkpoint: checkpoint
+        )
+        self.decoder = custom_decoder || TransformerDecoder.new(
+          num_decoder_layers,
+          dims,
+          num_heads,
+          mlp_dims: mlp_dims,
+          dropout: dropout,
+          activation: activation,
+          norm_first: norm_first,
+          checkpoint: checkpoint
+        )
+      end
+
+      def call(src, tgt, src_mask, tgt_mask, memory_mask)
+        memory = encoder.call(src, src_mask)
+        decoder.call(tgt, memory, tgt_mask, memory_mask)
+      end
+    end
+  end
+end

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

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+module MLX
+  module NN
+    class << self
+      def upsample_nearest(x, scale_factor)
+        dims = x.ndim - 2
+        if dims != scale_factor.length
+          raise ArgumentError, "A scale needs to be provided for each spatial dimension"
+        end
+
+        result = x
+        x.shape[1...-1].zip(scale_factor).each_with_index do |(n, scale), i|
+          indices = nearest_indices(n, scale)
+          result = MLX::Core.take(result, indices, 1 + i)
+        end
+        result
+      end
+
+      def upsample_linear(x, scale_factor, align_corners: false)
+        _ = align_corners
+        upsample_nearest(x, scale_factor)
+      end
+
+      def upsample_cubic(x, scale_factor, align_corners: false)
+        _ = align_corners
+        upsample_nearest(x, scale_factor)
+      end
+
+      private
+
+      def nearest_indices(n, scale)
+        m = (scale * n).to_i
+        if m <= 0
+          raise ArgumentError, "scale_factor must produce a positive output size"
+        end
+
+        if m > n
+          indices = Array.new(m) do |i|
+            (((i + 0.5) * (n.to_f / m)) - 0.5).round
+          end
+        else
+          indices = Array.new(m) do |i|
+            (i * (n.to_f / m)).floor
+          end
+        end
+        indices.map { |idx| [[idx, 0].max, n - 1].min }
+      end
+    end
+
+    class Upsample < Module
+      def initialize(scale_factor:, mode: "nearest", align_corners: false)
+        super()
+        unless %w[nearest linear cubic].include?(mode)
+          raise ArgumentError, "[Upsample] Got unsupported upsampling algorithm: #{mode}"
+        end
+
+        @scale_factor = if scale_factor.is_a?(Array)
+          scale_factor.map(&:to_f)
+        else
+          scale_factor.to_f
+        end
+        @mode = mode
+        @align_corners = align_corners
+      end
+
+      def call(x)
+        dims = x.ndim - 2
+        if dims <= 0
+          raise ArgumentError,
+                "[Upsample] The input should have at least 1 spatial dimension which means it should be at least 3D but #{x.ndim}D was provided"
+        end
+
+        scale = @scale_factor
+        if scale.is_a?(Array)
+          if scale.length != dims
+            raise ArgumentError,
+                  "[Upsample] One scale per spatial dimension is required but scale_factor=#{scale} and the number of spatial dimensions were #{dims}"
+          end
+        else
+          scale = Array.new(dims, scale)
+        end
+
+        case @mode
+        when "nearest"
+          MLX::NN.upsample_nearest(x, scale)
+        when "linear"
+          MLX::NN.upsample_linear(x, scale, align_corners: @align_corners)
+        when "cubic"
+          MLX::NN.upsample_cubic(x, scale, align_corners: @align_corners)
+        else
+          raise ArgumentError, "Unknown interpolation mode: #{@mode}"
+        end
+      end
+    end
+  end
+end

data/lib/mlx/nn/layers.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+require_relative "layers/activations"
+require_relative "layers/base"
+require_relative "layers/containers"
+require_relative "layers/convolution"
+require_relative "layers/convolution_transpose"
+require_relative "layers/distributed"
+require_relative "layers/dropout"
+require_relative "layers/embedding"
+require_relative "layers/linear"
+require_relative "layers/normalization"
+require_relative "layers/pooling"
+require_relative "layers/positional_encoding"
+require_relative "layers/quantized"
+require_relative "layers/recurrent"
+require_relative "layers/transformer"
+require_relative "layers/upsample"

data/lib/mlx/nn/losses.rb

@@ -0,0 +1,251 @@
+# frozen_string_literal: true
+
+module MLX
+  module NN
+    module Losses
+      module_function
+
+      VALID_REDUCTIONS = %w[none mean sum].freeze
+
+      def reduction(value, mode)
+        mode_name = (mode || "none").to_s
+        unless VALID_REDUCTIONS.include?(mode_name)
+          raise ArgumentError, "Invalid reduction. Must be one of #{VALID_REDUCTIONS}."
+        end
+
+        return MLX::Core.mean(value) if mode_name == "mean"
+        return MLX::Core.sum(value) if mode_name == "sum"
+
+        value
+      end
+
+      def cross_entropy(
+        logits,
+        targets,
+        weights: nil,
+        axis: -1,
+        label_smoothing: 0.0,
+        reduction: "none"
+      )
+        if label_smoothing < 0 || label_smoothing >= 1
+          raise ArgumentError, "Label smoothing must in [0, 1), got #{label_smoothing}."
+        end
+
+        targets_as_probs = targets.ndim == logits.ndim
+        dropped_shape = drop_dim(logits.shape, axis)
+        if (targets_as_probs && targets.shape != logits.shape) || (!targets_as_probs && targets.shape != dropped_shape)
+          raise ArgumentError, "Targets shape #{targets.shape} does not match logits shape #{logits.shape}."
+        end
+
+        score = if targets_as_probs
+          MLX::Core.sum(MLX::Core.multiply(logits, targets), axis)
+        else
+          idx = MLX::Core.expand_dims(targets, axis)
+          MLX::Core.squeeze(MLX::Core.take_along_axis(logits, idx, axis), axis)
+        end
+
+        logsumexp_logits = MLX::Core.logsumexp(logits, axis)
+        loss = if label_smoothing > 0
+          adjusted_score = MLX::Core.multiply(score, 1 - label_smoothing)
+          mean_logits = MLX::Core.mean(logits, axis)
+          smoothed_loss = MLX::Core.multiply(mean_logits, -label_smoothing)
+          MLX::Core.add(MLX::Core.subtract(logsumexp_logits, adjusted_score), smoothed_loss)
+        else
+          MLX::Core.subtract(logsumexp_logits, score)
+        end
+
+        unless weights.nil?
+          if weights.shape != loss.shape
+            raise ArgumentError,
+                  "Weights with shape #{weights.shape} is not the same as output loss with shape #{loss.shape}."
+          end
+          loss = MLX::Core.multiply(loss, weights)
+        end
+
+        reduction(loss, reduction)
+      end
+
+      def binary_cross_entropy(
+        inputs,
+        targets,
+        weights: nil,
+        with_logits: true,
+        reduction: "mean"
+      )
+        if inputs.shape != targets.shape
+          raise ArgumentError, "Inputs shape #{inputs.shape} does not match targets shape #{targets.shape}."
+        end
+
+        loss = if with_logits
+          MLX::Core.subtract(MLX::Core.logaddexp(0.0, inputs), MLX::Core.multiply(inputs, targets))
+        else
+          log_inputs_clip = MLX::Core.clip(MLX::Core.log(inputs), -100, nil)
+          log_inputs_inv_clip = MLX::Core.clip(
+            MLX::Core.log(MLX::Core.subtract(1.0, inputs)),
+            -100,
+            nil
+          )
+          weighted_terms = MLX::Core.add(
+            MLX::Core.multiply(targets, log_inputs_clip),
+            MLX::Core.multiply(MLX::Core.subtract(1.0, targets), log_inputs_inv_clip)
+          )
+          MLX::Core.multiply(weighted_terms, -1.0)
+        end
+
+        unless weights.nil?
+          if weights.shape != loss.shape
+            raise ArgumentError,
+                  "Weights with shape #{weights.shape} is not the same as output loss with shape #{loss.shape}."
+          end
+          loss = MLX::Core.multiply(loss, weights)
+        end
+
+        reduction(loss, reduction)
+      end
+
+      def l1_loss(predictions, targets, reduction: "mean")
+        if predictions.shape != targets.shape
+          raise ArgumentError, "Predictions shape #{predictions.shape} does not match targets shape #{targets.shape}."
+        end
+
+        reduction(MLX::Core.abs(MLX::Core.subtract(predictions, targets)), reduction)
+      end
+
+      def mse_loss(predictions, targets, reduction: "mean")
+        if predictions.shape != targets.shape
+          raise ArgumentError, "Predictions shape #{predictions.shape} does not match targets shape #{targets.shape}."
+        end
+
+        diff = MLX::Core.subtract(predictions, targets)
+        reduction(MLX::Core.square(diff), reduction)
+      end
+
+      def nll_loss(inputs, targets, axis: -1, reduction: "none")
+        idx = MLX::Core.expand_dims(targets, -1)
+        selected = MLX::Core.take_along_axis(inputs, idx, axis)
+        loss = MLX::Core.multiply(MLX::Core.squeeze(selected, -1), -1.0)
+        reduction(loss, reduction)
+      end
+
+      def gaussian_nll_loss(inputs, targets, vars, full: false, eps: 1e-6, reduction: "mean")
+        if inputs.shape != targets.shape
+          raise ArgumentError, "Inputs shape #{inputs.shape} does not match targets shape #{targets.shape}."
+        end
+        if inputs.shape != vars.shape
+          raise ArgumentError, "Inputs shape #{inputs.shape} does not match vars shape #{vars.shape}."
+        end
+
+        vars = MLX::Core.maximum(vars, eps)
+        squared_error = MLX::Core.square(MLX::Core.subtract(targets, inputs))
+        base = MLX::Core.add(MLX::Core.log(vars), MLX::Core.divide(squared_error, vars))
+        loss = MLX::Core.multiply(base, 0.5)
+        if full
+          loss = MLX::Core.add(loss, 0.5 * Math.log(2 * Math::PI))
+        end
+
+        reduction(loss, reduction)
+      end
+
+      def kl_div_loss(inputs, targets, axis: -1, reduction: "none")
+        diff = MLX::Core.subtract(targets, inputs)
+        loss = MLX::Core.sum(MLX::Core.multiply(MLX::Core.exp(targets), diff), axis)
+        reduction(loss, reduction)
+      end
+
+      def smooth_l1_loss(predictions, targets, beta: 1.0, reduction: "mean")
+        if predictions.shape != targets.shape
+          raise ArgumentError, "Predictions shape #{predictions.shape} does not match targets shape #{targets.shape}."
+        end
+
+        diff = MLX::Core.abs(MLX::Core.subtract(predictions, targets))
+        loss = MLX::Core.where(
+          MLX::Core.less(diff, beta),
+          MLX::Core.divide(MLX::Core.multiply(MLX::Core.square(diff), 0.5), beta),
+          MLX::Core.subtract(MLX::Core.abs(diff), 0.5 * beta)
+        )
+        reduction(loss, reduction)
+      end
+
+      def triplet_loss(
+        anchors,
+        positives,
+        negatives,
+        axis: -1,
+        p: 2,
+        margin: 1.0,
+        eps: 1e-6,
+        reduction: "none"
+      )
+        ap = MLX::Core.subtract(anchors, positives)
+        an = MLX::Core.subtract(anchors, negatives)
+        dist_ap = MLX::Core.sqrt(MLX::Core.add(MLX::Core.sum(MLX::Core.power(ap, p), axis), eps))
+        dist_an = MLX::Core.sqrt(MLX::Core.add(MLX::Core.sum(MLX::Core.power(an, p), axis), eps))
+        loss = MLX::Core.maximum(MLX::Core.add(MLX::Core.subtract(dist_ap, dist_an), margin), 0.0)
+        reduction(loss, reduction)
+      end
+
+      def hinge_loss(inputs, targets, reduction: "none")
+        loss = MLX::Core.maximum(MLX::Core.subtract(1.0, MLX::Core.multiply(inputs, targets)), 0.0)
+        reduction(loss, reduction)
+      end
+
+      def huber_loss(inputs, targets, delta: 1.0, reduction: "none")
+        errors = MLX::Core.subtract(inputs, targets)
+        abs_errors = MLX::Core.abs(errors)
+        quadratic = MLX::Core.minimum(abs_errors, delta)
+        linear = MLX::Core.subtract(abs_errors, quadratic)
+        loss = MLX::Core.add(
+          MLX::Core.multiply(MLX::Core.square(quadratic), 0.5),
+          MLX::Core.multiply(linear, delta)
+        )
+        reduction(loss, reduction)
+      end
+
+      def log_cosh_loss(inputs, targets, reduction: "none")
+        errors = MLX::Core.subtract(inputs, targets)
+        loss = MLX::Core.subtract(MLX::Core.logaddexp(errors, MLX::Core.multiply(errors, -1.0)), Math.log(2))
+        reduction(loss, reduction)
+      end
+
+      def cosine_similarity_loss(x1, x2, axis: 1, eps: 1e-8, reduction: "none")
+        x1_norm = MLX::Core.sqrt(MLX::Core.sum(MLX::Core.square(x1), axis))
+        x2_norm = MLX::Core.sqrt(MLX::Core.sum(MLX::Core.square(x2), axis))
+        dot = MLX::Core.sum(MLX::Core.multiply(x1, x2), axis)
+        denom = MLX::Core.maximum(MLX::Core.multiply(x1_norm, x2_norm), eps)
+        loss = MLX::Core.divide(dot, denom)
+        reduction(loss, reduction)
+      end
+
+      def margin_ranking_loss(inputs1, inputs2, targets, margin: 0.0, reduction: "none")
+        unless inputs1.shape == inputs2.shape && inputs1.shape == targets.shape
+          raise ArgumentError,
+                "The shapes of the arguments do not match. The provided shapes are " \
+                "inputs1.shape=#{inputs1.shape}, inputs2.shape=#{inputs2.shape}, and " \
+                "targets.shape=#{targets.shape}."
+        end
+
+        differences = MLX::Core.subtract(inputs1, inputs2)
+        term = MLX::Core.add(MLX::Core.multiply(MLX::Core.multiply(targets, differences), -1.0), margin)
+        loss = MLX::Core.maximum(0.0, term)
+        reduction(loss, reduction)
+      end
+
+      def drop_dim(shape, axis)
+        axis_index = axis.negative? ? axis + shape.length : axis
+        copied = shape.dup
+        copied.delete_at(axis_index)
+        copied
+      end
+    end
+
+    class << self
+      %i[
+        cross_entropy binary_cross_entropy l1_loss mse_loss nll_loss gaussian_nll_loss
+        kl_div_loss smooth_l1_loss triplet_loss hinge_loss huber_loss log_cosh_loss
+        cosine_similarity_loss margin_ranking_loss
+      ].each do |name|
+        define_method(name) { |*args, **kwargs| Losses.public_send(name, *args, **kwargs) }
+      end
+    end
+  end
+end