Use u8 tensors for masks. (#273)

huggingface · Jul 29, 2023 · 4bf2ebf · 4bf2ebf
1 parent 50d8273
commit 4bf2ebf
Show file tree

Hide file tree

Showing 6 changed files with 6 additions and 10 deletions.
diff --git a/candle-examples/examples/bigcode/model.rs b/candle-examples/examples/bigcode/model.rs
@@ -24,7 +24,7 @@ fn layer_norm(size: usize, eps: f64, vb: VarBuilder) -> Result<LayerNorm> {
 
 fn make_causal_mask(t: usize, device: &Device) -> Result<Tensor> {
     let mask: Vec<_> = (0..t)
-        .flat_map(|i| (0..t).map(move |j| u32::from(j <= i)))
+        .flat_map(|i| (0..t).map(move |j| u8::from(j <= i)))
         .collect();
     let mask = Tensor::from_slice(&mask, (t, t), device)?;
     Ok(mask)

diff --git a/candle-examples/examples/falcon/model.rs b/candle-examples/examples/falcon/model.rs
@@ -424,7 +424,7 @@ pub struct Falcon {
 
 fn make_causal_mask(t: usize) -> Result<Tensor> {
     let mask: Vec<_> = (0..t)
-        .flat_map(|i| (0..t).map(move |j| u32::from(j > i)))
+        .flat_map(|i| (0..t).map(move |j| u8::from(j > i)))
         .collect();
     let mask = Tensor::from_slice(&mask, (t, t), &Device::Cpu)?;
     Ok(mask)

diff --git a/candle-examples/examples/llama/model.rs b/candle-examples/examples/llama/model.rs
@@ -91,9 +91,8 @@ impl Cache {
         if let Some(mask) = masks.get(&t) {
             Ok(mask.clone())
         } else {
-            // TODO: If we support bool or u8 tensors, this would be better.
             let mask: Vec<_> = (0..t)
-                .flat_map(|i| (0..t).map(move |j| u32::from(j > i)))
+                .flat_map(|i| (0..t).map(move |j| u8::from(j > i)))
                 .collect();
             let mask = Tensor::from_slice(&mask, (t, t), &self.device)?;
             masks.insert(t, mask.clone());

diff --git a/candle-examples/examples/llama2-c/model.rs b/candle-examples/examples/llama2-c/model.rs
@@ -47,9 +47,8 @@ impl Cache {
         if let Some(mask) = masks.get(&t) {
             Ok(mask.clone())
         } else {
-            // TODO: If we support bool or u8 tensors, this would be better.
             let mask: Vec<_> = (0..t)
-                .flat_map(|i| (0..t).map(move |j| u32::from(j > i)))
+                .flat_map(|i| (0..t).map(move |j| u8::from(j > i)))
                 .collect();
             let mask = Tensor::from_slice(&mask, (t, t), &self.device)?;
             masks.insert(t, mask.clone());

diff --git a/candle-examples/examples/llama_multiprocess/model.rs b/candle-examples/examples/llama_multiprocess/model.rs
@@ -179,9 +179,8 @@ impl Cache {
         if let Some(mask) = masks.get(&t) {
             Ok(mask.clone())
         } else {
-            // TODO: If we support bool or u8 tensors, this would be better.
             let mask: Vec<_> = (0..t)
-                .flat_map(|i| (0..t).map(move |j| u32::from(j > i)))
+                .flat_map(|i| (0..t).map(move |j| u8::from(j > i)))
                 .collect();
             let mask = Tensor::from_slice(&mask, (t, t), &self.device)?;
             masks.insert(t, mask.clone());

diff --git a/candle-wasm-examples/llama2-c/src/model.rs b/candle-wasm-examples/llama2-c/src/model.rs
@@ -47,9 +47,8 @@ impl Cache {
         if let Some(mask) = masks.get(&t) {
             Ok(mask.clone())
         } else {
-            // TODO: If we support bool or u8 tensors, this would be better.
             let mask: Vec<_> = (0..t)
-                .flat_map(|i| (0..t).map(move |j| u32::from(j > i)))
+                .flat_map(|i| (0..t).map(move |j| u8::from(j > i)))
                 .collect();
             let mask = Tensor::from_slice(&mask, (t, t), &self.device)?;
             masks.insert(t, mask.clone());