fix #231 -- use Embedding.asLinear (#232)

Author: davidkoski

Libraries/MLXLLM/Models/Cohere.swift

@@ -163,7 +163,7 @@ public class CohereModel: Module, LLMModel, KVCacheDimensionProvider {
 
     public func callAsFunction(_ inputs: MLXArray, cache: [KVCache]?) -> MLXArray {
         var out = model(inputs, cache: cache)
-        out = matmul(out, model.embedTokens.weight.T)
+        out = model.embedTokens.asLinear(out)
         out = out * self.logitScale
         return out
     }

Libraries/MLXLLM/Models/OpenELM.swift

@@ -199,7 +199,7 @@ public class OpenELMModel: Module, LLMModel, KVCacheDimensionProvider {
         if let lmHead {
             out = lmHead(out)
         } else {
-            out = matmul(out, transformer.embedTokens.weight.T)
+            out = transformer.embedTokens.asLinear(out)
         }
 
         return out

Libraries/MLXLLM/Models/Starcoder2.swift

@@ -173,7 +173,7 @@ public class Starcoder2Model: Module, LLMModel, KVCacheDimensionProvider {
         if !tieWordEmbeddings {
             return lmHead(out)
         } else {
-            out = matmul(out, model.embedTokens.weight.T)
+            out = model.embedTokens.asLinear(out)
             return out
         }
     }