From 95abbc975b54b4a857faf341b17c3612509c7e8a Mon Sep 17 00:00:00 2001
From: Lukas Kreussel <65088241+LLukas22@users.noreply.github.com>
Date: Thu, 3 Aug 2023 12:11:30 +0200
Subject: [PATCH] remove `repeat` from falcon

---
 crates/models/falcon/src/lib.rs | 27 ++++++---------------------
 1 file changed, 6 insertions(+), 21 deletions(-)

diff --git a/crates/models/falcon/src/lib.rs b/crates/models/falcon/src/lib.rs
index 04a0b877..6e2a6833 100644
--- a/crates/models/falcon/src/lib.rs
+++ b/crates/models/falcon/src/lib.rs
@@ -171,11 +171,8 @@ impl KnownModel for Falcon {
                 // self-attention
                 layernorm_output = ctx0.op_norm(&input_layer);
                 layernorm_output = ctx0.op_add(
-                    &ctx0.op_mul(
-                        &ctx0.op_repeat(&self.layers[il].input_layernorm, &layernorm_output),
-                        &layernorm_output,
-                    ),
-                    &ctx0.op_repeat(&self.layers[il].input_layernorm_b, &layernorm_output),
+                    &ctx0.op_mul(&layernorm_output, &self.layers[il].input_layernorm),
+                    &self.layers[il].input_layernorm_b,
                 );
 
                 if n_head_kv == 1 {
@@ -185,17 +182,8 @@ impl KnownModel for Falcon {
                     // Falcon-40B only
                     current = ctx0.op_norm(&input_layer);
                     current = ctx0.op_add(
-                        &ctx0.op_mul(
-                            &ctx0.op_repeat(
-                                self.layers[il].attention_norm.as_ref().unwrap(),
-                                &current,
-                            ),
-                            &current,
-                        ),
-                        &ctx0.op_repeat(
-                            self.layers[il].attention_norm_b.as_ref().unwrap(),
-                            &current,
-                        ),
+                        &ctx0.op_mul(&current, self.layers[il].attention_norm.as_ref().unwrap()),
+                        self.layers[il].attention_norm_b.as_ref().unwrap(),
                     );
                 }
 
@@ -327,11 +315,8 @@ impl KnownModel for Falcon {
             input_layer = ctx0.op_norm(&input_layer);
 
             input_layer = ctx0.op_add(
-                &ctx0.op_mul(
-                    &ctx0.op_repeat(&self.output_norm, &input_layer),
-                    &input_layer,
-                ),
-                &ctx0.op_repeat(&self.output_norm_b, &input_layer),
+                &ctx0.op_mul(&input_layer, &self.output_norm),
+                &self.output_norm_b,
             );
 
             let embeddings_tensor: ggml::Tensor = input_layer.share();