mlx-ruby-lm 0.30.7.1

data/lib/mlx_lm/models/olmo2.rb

@@ -0,0 +1,169 @@
+module MlxLm
+  module Models
+    module OLMo2
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "olmo2"
+        field :hidden_size, default: 4096
+        field :num_hidden_layers, default: 32
+        field :num_attention_heads, default: 32
+        field :num_key_value_heads, default: nil
+        field :intermediate_size, default: 11008
+        field :vocab_size, default: 50304
+        field :rms_norm_eps, default: 1e-6
+        field :rope_theta, default: 10000.0
+        field :rope_traditional, default: false
+        field :rope_scaling, default: nil
+        field :attention_bias, default: false
+        field :mlp_bias, default: false
+        field :tie_word_embeddings, default: true
+        field :head_dim, default: nil
+        field :max_position_embeddings, default: nil
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+          @head_dim ||= @hidden_size / @num_attention_heads
+        end
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(args)
+          super()
+          dim = args.hidden_size
+          @n_heads = args.num_attention_heads
+          @n_kv_heads = args.num_key_value_heads
+          @head_dim = args.head_dim
+          @scale = @head_dim**(-0.5)
+
+          bias = args.attention_bias
+          self.q_proj = MLX::NN::Linear.new(dim, @n_heads * @head_dim, bias: bias)
+          self.k_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: bias)
+          self.v_proj = MLX::NN::Linear.new(dim, @n_kv_heads * @head_dim, bias: bias)
+          self.o_proj = MLX::NN::Linear.new(@n_heads * @head_dim, dim, bias: bias)
+
+          # OLMo2: Q and K normalization
+          self.q_norm = MLX::NN::RMSNorm.new(@head_dim, eps: args.rms_norm_eps)
+          self.k_norm = MLX::NN::RMSNorm.new(@head_dim, eps: args.rms_norm_eps)
+
+          self.rope = MLX::NN::RoPE.new(
+            @head_dim,
+            traditional: args.rope_traditional,
+            base: args.rope_theta
+          )
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, _d = x.shape
+
+          queries = q_proj.call(x).reshape([b, l, @n_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = k_proj.call(x).reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = v_proj.call(x).reshape([b, l, @n_kv_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          # Apply Q/K normalization
+          queries = q_norm.call(queries)
+          keys = k_norm.call(keys)
+
+          if cache
+            queries = rope.call(queries, offset: cache.offset)
+            keys = rope.call(keys, offset: cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+          else
+            queries = rope.call(queries)
+            keys = rope.call(keys)
+          end
+
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, @n_heads * @head_dim])
+          o_proj.call(output)
+        end
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(dim, hidden_dim, bias: false)
+          super()
+          self.gate_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: bias)
+          self.down_proj = MLX::NN::Linear.new(hidden_dim, dim, bias: bias)
+          self.up_proj = MLX::NN::Linear.new(dim, hidden_dim, bias: bias)
+        end
+
+        def call(x)
+          down_proj.call(MLX::NN.silu(gate_proj.call(x)) * up_proj.call(x))
+        end
+      end
+
+      class TransformerBlock < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.self_attn = Attention.new(args)
+          self.mlp = MLP.new(args.hidden_size, args.intermediate_size, bias: args.mlp_bias)
+          self.input_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+          self.post_attention_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          r = self_attn.call(input_layernorm.call(x), mask: mask, cache: cache)
+          h = x + r
+          r = mlp.call(post_attention_layernorm.call(h))
+          h + r
+        end
+      end
+
+      class OLMo2Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.embed_tokens = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.layers = Array.new(args.num_hidden_layers) { TransformerBlock.new(args) }
+          self.norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(inputs, cache: nil)
+          h = embed_tokens.call(inputs)
+          layer_cache = cache || [nil] * layers.length
+
+          mask = nil
+          mask = "causal" if h.shape[1] > 1
+
+          layers.each_with_index do |layer, i|
+            h = layer.call(h, mask: mask, cache: layer_cache[i])
+          end
+
+          norm.call(h)
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model = OLMo2Model.new(args)
+          unless args.tie_word_embeddings
+            self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+          end
+        end
+
+        def call(inputs, cache: nil)
+          out = model.call(inputs, cache: cache)
+          if @args.tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def sanitize(weights)
+          result = weights.reject { |k, _| k.include?("self_attn.rotary_emb.inv_freq") }
+          result.delete("lm_head.weight") if @args.tie_word_embeddings
+          result
+        end
+
+        def layers
+          model.layers
+        end
+      end
+
+      Models.register("olmo2", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/olmo3.rb

@@ -0,0 +1,254 @@
+module MlxLm
+  module Models
+    module OLMo3
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "olmo3"
+        field :hidden_size
+        field :num_hidden_layers
+        field :intermediate_size
+        field :num_attention_heads
+        field :rms_norm_eps
+        field :vocab_size
+        field :max_position_embeddings
+        field :sliding_window
+        field :rope_theta
+        field :attention_bias, default: false
+        field :layer_types, default: nil
+        field :num_key_value_heads, default: nil
+        field :head_dim, default: nil
+        field :rope_scaling, default: nil
+        field :tie_word_embeddings, default: false
+
+        def initialize(**kwargs)
+          super
+          @num_key_value_heads ||= @num_attention_heads
+          @head_dim ||= @hidden_size / @num_attention_heads
+          @layer_types ||= Array.new(@num_hidden_layers) do |i|
+            ((i + 1) % 4).zero? ? "full_attention" : "sliding_attention"
+          end
+        end
+      end
+
+      class Olmo3Attention < MLX::NN::Module
+        def initialize(args, layer_idx:)
+          super()
+          @num_attention_heads = args.num_attention_heads
+          @num_key_value_heads = args.num_key_value_heads
+          @head_dim = args.head_dim
+          @scale = @head_dim**(-0.5)
+
+          self.q_proj = MLX::NN::Linear.new(
+            args.hidden_size,
+            @num_attention_heads * @head_dim,
+            bias: args.attention_bias
+          )
+          self.k_proj = MLX::NN::Linear.new(
+            args.hidden_size,
+            @num_key_value_heads * @head_dim,
+            bias: args.attention_bias
+          )
+          self.v_proj = MLX::NN::Linear.new(
+            args.hidden_size,
+            @num_key_value_heads * @head_dim,
+            bias: args.attention_bias
+          )
+          self.o_proj = MLX::NN::Linear.new(
+            @num_attention_heads * @head_dim,
+            args.hidden_size,
+            bias: args.attention_bias
+          )
+
+          self.q_norm = MLX::NN::RMSNorm.new(
+            @num_attention_heads * @head_dim,
+            eps: args.rms_norm_eps
+          )
+          self.k_norm = MLX::NN::RMSNorm.new(
+            @num_key_value_heads * @head_dim,
+            eps: args.rms_norm_eps
+          )
+
+          if args.layer_types[layer_idx] != "full_attention"
+            self.rope = MLX::NN::RoPE.new(@head_dim, traditional: false, base: args.rope_theta)
+          else
+            self.rope = MlxLm::Models.initialize_rope(
+              @head_dim,
+              args.rope_theta,
+              false,
+              args.rope_scaling,
+              max_position_embeddings: args.max_position_embeddings
+            )
+          end
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, _d = x.shape
+
+          queries = q_norm.call(q_proj.call(x))
+          keys = k_norm.call(k_proj.call(x))
+          values = v_proj.call(x)
+
+          queries = queries.reshape([b, l, @num_attention_heads, @head_dim]).transpose([0, 2, 1, 3])
+          keys = keys.reshape([b, l, @num_key_value_heads, @head_dim]).transpose([0, 2, 1, 3])
+          values = values.reshape([b, l, @num_key_value_heads, @head_dim]).transpose([0, 2, 1, 3])
+
+          if cache
+            queries = rope.call(queries, offset: cache.offset)
+            keys = rope.call(keys, offset: cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+          else
+            queries = rope.call(queries)
+            keys = rope.call(keys)
+          end
+
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, @num_attention_heads * @head_dim])
+          o_proj.call(output)
+        end
+      end
+
+      class Olmo3MLP < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.gate_proj = MLX::NN::Linear.new(args.hidden_size, args.intermediate_size, bias: false)
+          self.down_proj = MLX::NN::Linear.new(args.intermediate_size, args.hidden_size, bias: false)
+          self.up_proj = MLX::NN::Linear.new(args.hidden_size, args.intermediate_size, bias: false)
+        end
+
+        def call(x)
+          down_proj.call(Activations.swiglu(gate_proj.call(x), up_proj.call(x)))
+        end
+      end
+
+      class Olmo3DecoderLayer < MLX::NN::Module
+        def initialize(args, layer_idx:)
+          super()
+          self.self_attn = Olmo3Attention.new(args, layer_idx: layer_idx)
+          self.mlp = Olmo3MLP.new(args)
+          self.post_attention_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+          self.post_feedforward_layernorm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          r = post_attention_layernorm.call(self_attn.call(x, mask: mask, cache: cache))
+          h = x + r
+          r = post_feedforward_layernorm.call(mlp.call(h))
+          h + r
+        end
+      end
+
+      class Olmo3Model < MLX::NN::Module
+        attr_reader :layer_types
+
+        def initialize(args)
+          super()
+          @sliding_window = args.sliding_window
+          @layer_types = args.layer_types
+
+          self.embed_tokens = MLX::NN::Embedding.new(args.vocab_size, args.hidden_size)
+          self.layers = Array.new(args.num_hidden_layers) do |i|
+            Olmo3DecoderLayer.new(args, layer_idx: i)
+          end
+          self.norm = MLX::NN::RMSNorm.new(args.hidden_size, eps: args.rms_norm_eps)
+
+          self.swa_idx = @layer_types.index("sliding_attention") || 0
+          self.ga_idx = @layer_types.index("full_attention") || 0
+        end
+
+        def call(inputs, cache: nil)
+          h = embed_tokens.call(inputs)
+          layer_cache = cache || [nil] * layers.length
+
+          full_mask = _create_attention_mask(h, layer_cache[ga_idx])
+          sliding_window_mask = _create_attention_mask(
+            h,
+            layer_cache[swa_idx],
+            window_size: @sliding_window
+          )
+
+          layers.each_with_index do |layer, i|
+            mask = @layer_types[i] == "full_attention" ? full_mask : sliding_window_mask
+            h = layer.call(h, mask: mask, cache: layer_cache[i])
+          end
+
+          norm.call(h)
+        end
+
+        private
+
+        def _create_attention_mask(h, cache = nil, window_size: nil)
+          n = h.shape[1]
+          if cache && cache.respond_to?(:make_mask)
+            return cache.make_mask(n, window_size: window_size)
+          end
+
+          if window_size
+            offset = 0
+            if cache
+              offset = cache.offset
+              if cache.instance_variable_defined?(:@max_size)
+                max_size = cache.instance_variable_get(:@max_size)
+                offset = [max_size - 1, offset].min if max_size && max_size > 0
+              end
+            end
+            return _create_causal_mask(n, offset: offset, window_size: window_size) if offset + n > window_size
+          end
+          return nil if n == 1
+
+          "causal"
+        end
+
+        def _create_causal_mask(n, offset: 0, window_size: nil)
+          mx = MLX::Core
+          rinds = mx.arange(0, offset + n, 1, mx.int32).reshape([1, offset + n])
+          linds = mx.arange(offset, offset + n, 1, mx.int32).reshape([n, 1])
+
+          mask = mx.greater_equal(linds, rinds)
+          if window_size
+            mask = mx.logical_and(mask, mx.less(linds, mx.add(rinds, window_size)))
+          end
+          mask
+        end
+      end
+
+      class Model < MLX::NN::Module
+        attr_reader :args
+
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.model = Olmo3Model.new(args)
+          unless args.tie_word_embeddings
+            self.lm_head = MLX::NN::Linear.new(args.hidden_size, args.vocab_size, bias: false)
+          end
+        end
+
+        def call(inputs, cache: nil)
+          out = model.call(inputs, cache: cache)
+          if args.tie_word_embeddings
+            model.embed_tokens.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def layers
+          model.layers
+        end
+
+        def make_cache
+          model.layer_types.map do |layer_type|
+            if layer_type == "full_attention"
+              KVCache.new
+            else
+              RotatingKVCache.new(max_size: args.sliding_window)
+            end
+          end
+        end
+      end
+
+      Models.register("olmo3", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/olmoe.rb

@@ -0,0 +1,64 @@
+require_relative "olmo2"
+
+module MlxLm
+  module Models
+    module OLMoE
+      class ModelArgs < OLMo2::ModelArgs
+        field :model_type, default: "olmoe"
+        field :num_experts
+        field :num_experts_per_tok
+        field :norm_topk_prob, default: false
+      end
+
+      class Model < OLMo2::Model
+        def sanitize(weights)
+          result = super(weights)
+          rewrite_expert_weights(result)
+        end
+
+        private
+
+        def rewrite_expert_weights(weights)
+          return weights unless weights.key?("model.layers.0.mlp.experts.0.up_proj.weight")
+
+          mx = MLX::Core
+
+          layers.length.times do |layer_idx|
+            prefix = "model.layers.#{layer_idx}.mlp"
+            %w[up_proj down_proj gate_proj].each do |projection|
+              %w[weight scales biases].each do |param|
+                first_key = "#{prefix}.experts.0.#{projection}.#{param}"
+                next unless weights.key?(first_key)
+
+                expert_count = @args.num_experts || infer_expert_count(weights, prefix, projection, param)
+                next unless expert_count && expert_count.positive?
+
+                expert_keys = (0...expert_count).map do |expert_idx|
+                  "#{prefix}.experts.#{expert_idx}.#{projection}.#{param}"
+                end
+                next unless expert_keys.all? { |key| weights.key?(key) }
+
+                weights["#{prefix}.switch_mlp.#{projection}.#{param}"] = mx.stack(expert_keys.map { |key| weights.delete(key) })
+              end
+            end
+          end
+
+          weights
+        end
+
+        def infer_expert_count(weights, prefix, projection, param)
+          pattern = /\A#{Regexp.escape(prefix)}\.experts\.(\d+)\.#{projection}\.#{param}\z/
+          indices = weights.keys.filter_map do |key|
+            match = pattern.match(key)
+            match[1].to_i if match
+          end
+          return 0 if indices.empty?
+
+          indices.max + 1
+        end
+      end
+
+      Models.register("olmoe", Model, ModelArgs)
+    end
+  end
+end

data/lib/mlx_lm/models/openelm.rb

@@ -0,0 +1,208 @@
+module MlxLm
+  module Models
+    module OpenELM
+      module_function
+
+      def make_divisible(v, divisor = 8, min_value = nil)
+        min_value ||= divisor
+        rounded = ((v + (divisor.to_f / 2)).to_i / divisor) * divisor
+        new_v = [min_value, rounded].max
+        new_v += divisor if new_v < (0.9 * v)
+        new_v
+      end
+
+      class ModelArgs < BaseModelArgs
+        field :model_type, default: "openelm"
+        field :head_dim, default: 64
+        field :num_transformer_layers, default: 12
+        field :model_dim, default: 2048
+        field :vocab_size, default: 32_000
+        field :ffn_dim_divisor, default: 8
+        field :num_query_heads, default: [32]
+        field :num_kv_heads, default: []
+        field :ffn_multipliers, default: [1.0]
+        field :ffn_with_glu, default: true
+        field :normalize_qk_projections, default: true
+        field :share_input_output_layers, default: true
+        field :rms_norm_eps, default: 1e-6
+        field :rope_freq_constant, default: 10_000.0
+
+        def initialize(**kwargs)
+          super
+          @num_query_heads = normalize_schedule(@num_query_heads, @num_transformer_layers, 1, "num_query_heads").map(&:to_i)
+
+          if @num_kv_heads.nil? || Array(@num_kv_heads).empty?
+            @num_kv_heads = @num_query_heads.dup
+          else
+            @num_kv_heads = normalize_schedule(@num_kv_heads, @num_transformer_layers, @num_query_heads[0], "num_kv_heads").map(&:to_i)
+          end
+
+          @ffn_multipliers = normalize_schedule(@ffn_multipliers, @num_transformer_layers, 1.0, "ffn_multipliers").map(&:to_f)
+        end
+
+        private
+
+        def normalize_schedule(values, layers, fallback, field_name)
+          items = Array(values)
+          items = [fallback] if items.empty?
+          items = Array.new(layers, items[0]) if items.length == 1 && layers > 1
+
+          unless items.length == layers
+            raise ArgumentError, "#{field_name} must have #{layers} entries, got #{items.length}"
+          end
+
+          items
+        end
+      end
+
+      class Attention < MLX::NN::Module
+        def initialize(args, layer_id:)
+          super()
+          @head_dim = args.head_dim
+          @n_heads = args.num_query_heads[layer_id]
+          @n_kv_heads = args.num_kv_heads[layer_id]
+          @scale = @head_dim**(-0.5)
+          @normalize_qk_projections = args.normalize_qk_projections
+
+          op_size = (@n_heads + (2 * @n_kv_heads)) * @head_dim
+          self.qkv_proj = MLX::NN::Linear.new(args.model_dim, op_size, bias: false)
+          self.out_proj = MLX::NN::Linear.new(@n_heads * @head_dim, args.model_dim, bias: false)
+
+          if @normalize_qk_projections
+            self.q_norm = MLX::NN::RMSNorm.new(@head_dim, eps: args.rms_norm_eps)
+            self.k_norm = MLX::NN::RMSNorm.new(@head_dim, eps: args.rms_norm_eps)
+          end
+
+          self.rope = MLX::NN::RoPE.new(@head_dim, traditional: false, base: args.rope_freq_constant)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          mx = MLX::Core
+          b, l, _d = x.shape
+
+          qkv = qkv_proj.call(x)
+          qkv = qkv.reshape([b, l, @n_heads + (2 * @n_kv_heads), @head_dim]).transpose([0, 2, 1, 3])
+          queries, keys, values = mx.split(qkv, [@n_heads, @n_heads + @n_kv_heads], 1)
+
+          if @normalize_qk_projections
+            queries = q_norm.call(queries)
+            keys = k_norm.call(keys)
+          end
+
+          if cache
+            queries = rope.call(queries, offset: cache.offset)
+            keys = rope.call(keys, offset: cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+          else
+            queries = rope.call(queries)
+            keys = rope.call(keys)
+          end
+
+          output = mx.scaled_dot_product_attention(queries, keys, values, @scale, mask)
+          output = output.transpose([0, 2, 1, 3]).reshape([b, l, @n_heads * @head_dim])
+          out_proj.call(output)
+        end
+      end
+
+      class MLP < MLX::NN::Module
+        def initialize(args, layer_id:)
+          super()
+          @ffn_with_glu = args.ffn_with_glu
+          dim = args.model_dim
+          multiplier = args.ffn_multipliers[layer_id]
+          @intermediate_dim = OpenELM.make_divisible(multiplier * dim, args.ffn_dim_divisor).to_i
+
+          proj_1_dim = @ffn_with_glu ? (2 * @intermediate_dim) : @intermediate_dim
+          self.proj_1 = MLX::NN::Linear.new(dim, proj_1_dim, bias: false)
+          self.proj_2 = MLX::NN::Linear.new(@intermediate_dim, dim, bias: false)
+        end
+
+        def call(x)
+          x = proj_1.call(x)
+          x = if @ffn_with_glu
+            gate, value = MLX::Core.split(x, [@intermediate_dim], -1)
+            Activations.swiglu(gate, value)
+          else
+            MLX::NN.gelu_approx(x)
+          end
+          proj_2.call(x)
+        end
+      end
+
+      class TransformerBlock < MLX::NN::Module
+        def initialize(args, layer_id:)
+          super()
+          self.attn = Attention.new(args, layer_id: layer_id)
+          self.ffn = MLP.new(args, layer_id: layer_id)
+          self.attn_norm = MLX::NN::RMSNorm.new(args.model_dim, eps: args.rms_norm_eps)
+          self.ffn_norm = MLX::NN::RMSNorm.new(args.model_dim, eps: args.rms_norm_eps)
+        end
+
+        def call(x, mask: nil, cache: nil)
+          r = attn.call(attn_norm.call(x), mask: mask, cache: cache)
+          h = x + r
+          r = ffn.call(ffn_norm.call(h))
+          h + r
+        end
+      end
+
+      class OpenELMModel < MLX::NN::Module
+        def initialize(args)
+          super()
+          self.token_embeddings = MLX::NN::Embedding.new(args.vocab_size, args.model_dim)
+          self.layers = Array.new(args.num_transformer_layers) do |layer_id|
+            TransformerBlock.new(args, layer_id: layer_id)
+          end
+          self.norm = MLX::NN::RMSNorm.new(args.model_dim, eps: args.rms_norm_eps)
+        end
+
+        def call(inputs, cache: nil)
+          h = token_embeddings.call(inputs)
+          layer_cache = cache || [nil] * layers.length
+          mask = _create_attention_mask(h, layer_cache[0])
+
+          layers.each_with_index do |layer, i|
+            h = layer.call(h, mask: mask, cache: layer_cache[i])
+          end
+          norm.call(h)
+        end
+
+        private
+
+        def _create_attention_mask(h, cache)
+          return cache.make_mask(h.shape[1]) if cache && cache.respond_to?(:make_mask)
+          return nil if h.shape[1] == 1
+
+          "causal"
+        end
+      end
+
+      class Model < MLX::NN::Module
+        def initialize(args)
+          super()
+          @args = args
+          self.model_type = args.model_type
+          self.transformer = OpenELMModel.new(args)
+          unless args.share_input_output_layers
+            self.lm_head = MLX::NN::Linear.new(args.model_dim, args.vocab_size, bias: false)
+          end
+        end
+
+        def call(inputs, cache: nil)
+          out = transformer.call(inputs, cache: cache)
+          if @args.share_input_output_layers
+            transformer.token_embeddings.as_linear(out)
+          else
+            lm_head.call(out)
+          end
+        end
+
+        def layers
+          transformer.layers
+        end
+      end
+
+      Models.register("openelm", Model, ModelArgs)
+    end
+  end
+end