Fixed examples

examples/llama/prompts/merge_sort.txt

@@ -1,5 +1,5 @@
 <|begin_of_text|><|start_header_id|>system<|end_header_id|>
 
-You are a helpful AI assistant for travel tips and recommendations<|eot_id|><|start_header_id|>user<|end_header_id|>
+You are a helpful AI assistant<|eot_id|><|start_header_id|>user<|end_header_id|>
 
-What can you help me with?<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+Please write a python implementation of merge sort.<|eot_id|><|start_header_id|>assistant<|end_header_id|>

examples/llama/src/main.rs

@@ -21,7 +21,7 @@ use luminal::prelude::*;
 #[command(author, version, about, long_about = None)]
 pub struct CLIArgs {
     /// Number of tokens to generate
-    #[clap(short = 't', long = "gen_tokens", default_value = "128")]
+    #[clap(short = 't', long = "gen_tokens", default_value = "256")]
     gen_tokens: i32,
 
     /// Prompt for the model

src/hl_ops/matmul.rs

@@ -78,7 +78,6 @@ impl GraphTensor {
                 )
             }
         } else if self.shape.len() == 5 && rhs.shape.len() == 5 {
-            println!("5x5");
             // ABCDExABCEF -> ABCDF
             let (a, b, c, e, f) = rhs.dims5();
             let (_, _, _, d, _) = self.dims5();

src/shape/symbolic.rs

@@ -953,10 +953,10 @@ fn make_rules() -> Vec<Rewrite> {
         rewrite!("distribute-max"; "(* ?a (max ?b ?c))"        => "(max (* ?a ?b) (* ?a ?c))"),
         rewrite!("distribute-min"; "(* ?a (min ?b ?c))"        => "(min (* ?a ?b) (* ?a ?c))"),
         rewrite!("factor"    ; "(+ (* ?a ?b) (* ?a ?c))" => "(* ?a (+ ?b ?c))"),
-        // rewrite!("group-terms"; "(+ ?a ?a)" => "(* 2 ?a)"),
-        // rewrite!("distribute-mod"; "(* (% ?b ?c) ?a)" => "(% (* ?b ?a) (* ?c ?a))"),
-        // rewrite!("explicit-truncate"; "(* (/ ?a ?b) ?b)" => "(- ?a (% ?a ?b))"),
-        // rewrite!("mul-mod"; "(% (* ?a ?b) ?b)" => "0"),
+        rewrite!("group-terms"; "(+ ?a ?a)" => "(* 2 ?a)"),
+        rewrite!("distribute-mod"; "(* (% ?b ?c) ?a)" => "(% (* ?b ?a) (* ?c ?a))"),
+        rewrite!("explicit-truncate"; "(* (/ ?a ?b) ?b)" => "(- ?a (% ?a ?b))"),
+        rewrite!("mul-mod"; "(% (* ?a ?b) ?b)" => "0"),
         // rewrite!("mul-distribute"; "(* ?a (% (/ ?b ?c) ?d))" => "(% (/ ?b (* ?c ?a)) (* ?d ?a))"),
         // rewrite!("div-mod-mul"; "(% (/ ?a ?b) ?c)" => "(% ?a (* ?b ?c))"),
     ]