Merge branch 'main' of https://github.com/jafioti/luminal

jafioti · Apr 26, 2024 · 7357817 · 7357817
2 parents a7f8230 + 862af13
commit 7357817
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 23 deletions.
diff --git a/crates/luminal_metal/src/elementwise_fusion.rs b/crates/luminal_metal/src/elementwise_fusion.rs
@@ -407,7 +407,7 @@ impl<T: MetalFloat> Compiler for ElementwiseFusionCompiler<T> {
                             },
                         )
                         .0;
-                    if val_exp != true.into() {
+                    if val_exp != true {
                         *subexp = format!(
                             "(({} != 0) ? {subexp} : 0.0)",
                             expr_to_metal_string(val_exp)
@@ -875,8 +875,8 @@ mod tests {
                     .permute::<_, Axes4<0, 2, 1, 3>>();
 
                 // Rotary embed queries and keys
-                let queries = apply_rotary_embeddings_ggml(queries, PrevSeq::const_size().into());
-                let keys = apply_rotary_embeddings_ggml(keys, PrevSeq::const_size().into());
+                let queries = apply_rotary_embeddings_ggml(queries, PrevSeq::const_size().big());
+                let keys = apply_rotary_embeddings_ggml(keys, PrevSeq::const_size().big());
 
                 // Add KV cache
                 let (keys, values) = (

diff --git a/crates/luminal_nn/src/convolution.rs b/crates/luminal_nn/src/convolution.rs
@@ -23,18 +23,15 @@ impl<
     for Conv1D<CHANNELS_IN, CHANNELS_OUT, KERNEL, STRIDE, DILATION, CHANNELS_IN_TIMES_KERNEL>
 {
     fn initialize(cx: &mut Graph) -> Self {
-        let conv = Self {
-            weight: cx.named_tensor("Weight"),
-        };
-
         // Init weight as uniform(-1, 1)
         let mut rng = thread_rng();
-        conv.weight.set(
-            (0..(CHANNELS_IN * CHANNELS_OUT * KERNEL))
-                .map(|_| rng.gen_range(-1_f32..1_f32))
-                .collect::<Vec<_>>(),
-        );
-        conv
+        Self {
+            weight: cx.named_tensor("Weight").set(
+                (0..(CHANNELS_IN * CHANNELS_OUT * KERNEL))
+                    .map(|_| rng.gen_range(-1_f32..1_f32))
+                    .collect::<Vec<_>>(),
+            ),
+        }
     }
 }
 
@@ -118,18 +115,15 @@ impl<
     >
 {
     fn initialize(cx: &mut Graph) -> Self {
-        let conv = Self {
-            weight: cx.named_tensor("Weight"),
-        };
-
         // Init weight as uniform(-1, 1)
         let mut rng = thread_rng();
-        conv.weight.set(
-            (0..(CHANNELS_IN * CHANNELS_OUT * KERNELX * KERNELY))
-                .map(|_| rng.gen_range(-1_f32..1_f32))
-                .collect::<Vec<_>>(),
-        );
-        conv
+        Self {
+            weight: cx.named_tensor("Weight").set(
+                (0..(CHANNELS_IN * CHANNELS_OUT * KERNELX * KERNELY))
+                    .map(|_| rng.gen_range(-1_f32..1_f32))
+                    .collect::<Vec<_>>(),
+            ),
+        }
     }
 }