Update llm tests

.travis/cache_file.sh

@@ -12,9 +12,12 @@ cd $CACHEDIR
 for file in $@
 do
     mkdir -p $(dirname $file)
-    [ -e $file ] \
-        || wget --no-verbose https://s3.amazonaws.com/tract-ci-builds/tests/$file -O $file \
-        || aws s3 cp s3://tract-ci-builds/tests/$file $file
+    if [ ! -e $file ]
+    then
+        wget --no-verbose https://s3.amazonaws.com/tract-ci-builds/tests/$file -O $file.tmp \
+        || aws s3 cp s3://tract-ci-builds/tests/$file $file.tmp
+        mv $file.tmp $file
+    fi
 done
 
 exit 0

.travis/test-llm.sh

@@ -28,25 +28,44 @@ case $model in
         ;;
 esac
 
-case $q in
-    q40f16) approx=ultra;;
-    q40ef16) approx=ultra;;
-    f16f16) approx=ultra;;
-    q40f32) approx=very;;
-    q40ef32) approx=very;;
-    f32f32) approx=approximate;;
-esac
-
 nnef=llm/$generation/$id/$id.nnef.tgz
 
+scenarios="p0s100 p99s1"
+
+if [ "$model" != "phi-1_5" ]
+then
+    scenarios="p0s100 p99s1 p50s50"
+fi
+
 set -x
 $CACHE_FILE $nnef
-for t in p0s100 p50s50 p99s1 
+for t in $scenarios
 do
     npz=llm/$generation/$id/$id.$t.io.npz
     $CACHE_FILE $npz
+
+    case $q in
+        q40f16) approx="--approx ultra";;
+        q40ef16) approx="--approx ultra";;
+        f16f16) approx="--approx ultra";;
+        q40f32) approx="--approx very";;
+        q40ef32) approx="--approx very";;
+        f32f32) approx="--approx approximate";;
+    esac
+
+    case "$id.$t" in 
+        apple--OpenELM-270M-f16f16.p50s50) approx="--approx-custom 0.2,0.1,0.003";;
+        TinyLlama--TinyLlama_v1.1-f16f16.p0s100) approx="--approx-custom 0.2,0.1,0.001";;
+        TinyLlama--TinyLlama_v1.1-f16f16.p50s50) approx="--approx-custom 0.2,0.1,0.005";;
+        TinyLlama--TinyLlama_v1.1-f16f16.p99s1) approx="--approx-custom 0.2,0.1,0.004";;
+        TinyLlama--TinyLlama_v1.1-q40f16.p0s100) approx="--approx-custom 0.2,0.1,0.004";;
+        TinyLlama--TinyLlama_v1.1-q40f16.p99s1) approx="--approx-custom 0.2,0.1,0.002";;
+        TinyLlama--TinyLlama_v1.1-q40f16.p50s50) approx="--approx-custom 0.2,0.1,0.004";;
+    esac
+
+
     $TRACT_RUN -v --nnef-tract-core $MODELS/$nnef -O run \
         --input-from-npz $MODELS/$npz \
         --assert-output-bundle $MODELS/$npz \
-        --approx $approx --allow-float-casts
+        $approx --allow-float-casts
 done

cli/src/main.rs

@@ -421,6 +421,12 @@ fn assertions_options(command: clap::Command) -> clap::Command {
             .long("approx")
             .help("Approximation level used in assertions."),
             )
+        .arg(
+            Arg::new("approx-custom")
+            .takes_value(true)
+            .long("approx-custom")
+            .help("Approximation level used in assertions (atol, rtol, outlier ratio). 3 coma-separated floats."),
+            )
         .arg(
             Arg::new("assert-output")
             .takes_value(true)

cli/src/params.rs

@@ -1134,14 +1134,19 @@ impl Assertions {
                     .collect()
             });
         let allow_missing_outputs = sub.is_present("allow-missing-outputs");
-        let approximation = match sub.value_of("approx").unwrap() {
-            "exact" => Approximation::Exact,
-            "close" => Approximation::Close,
-            "approximate" => Approximation::Approximate,
-            "very" => Approximation::VeryApproximate,
-            "super" => Approximation::SuperApproximate,
-            "ultra" => Approximation::UltraApproximate,
-            _ => panic!(),
+        let approximation = if let Some(custom) = sub.value_of("approx-custom") {
+            let Some((atol, rtol, approx)) = custom.split(",").collect_tuple() else { bail!("Can't parse approx custom. It should look like 0.001,0.002,0.003") };
+            Approximation::Custom(atol.parse()?, rtol.parse()?, approx.parse()?)
+        } else {
+            match sub.value_of("approx").unwrap() {
+                "exact" => Approximation::Exact,
+                "close" => Approximation::Close,
+                "approximate" => Approximation::Approximate,
+                "very" => Approximation::VeryApproximate,
+                "super" => Approximation::SuperApproximate,
+                "ultra" => Approximation::UltraApproximate,
+                _ => panic!(),
+            }
         };
         Ok(Assertions {
             assert_outputs,

data/src/tensor.rs

@@ -20,7 +20,7 @@ use std::sync::Arc;
 pub mod litteral;
 pub mod view;
 
-#[derive(Copy, Clone, Default, PartialEq, Eq, Debug)]
+#[derive(Copy, Clone, Default, Debug)]
 pub enum Approximation {
     Exact,
     #[default]
@@ -29,8 +29,22 @@ pub enum Approximation {
     VeryApproximate,
     SuperApproximate,
     UltraApproximate,
+    Custom(f32, f32, f32),
 }
 
+impl PartialEq for Approximation {
+    fn eq(&self, other: &Self) -> bool {
+        use Approximation::Custom;
+        if let (Custom(aa, ar, ao), Custom(ba, br, bo)) = (self, other) {
+            aa == ba && ar == br && bo == ao
+        } else {
+            std::mem::discriminant(self) == std::mem::discriminant(other)
+        }
+    }
+}
+
+impl Eq for Approximation {}
+
 impl From<bool> for Approximation {
     fn from(b: bool) -> Self {
         if b {
@@ -54,6 +68,7 @@ impl Approximation {
             (VeryApproximate, _) => (5e-2, 1e-2, 0.0),
             (SuperApproximate, _) => (0.1, 0.05, 0.0001),
             (UltraApproximate, _) => (0.2, 0.1, 0.0005),
+            (Custom(atol, rtol, out), _) => (*atol as _, *rtol as _, *out as _),
         }
     }
 }