Add masked_fill under Tensor

candle-core/src/tensor.rs

@@ -2454,6 +2454,13 @@ impl Tensor {
     pub fn broadcast_pow(&self, rhs: &Tensor) -> Result<Self> {
         rhs.broadcast_mul(&self.log()?)?.exp()
     }
+
+    pub fn masked_fill(&self, rhs: &Tensor, value: f32) -> Result<Self> {
+        rhs.where_cond(
+            &Tensor::new(value, self.device())?.broadcast_as(rhs.shape().dims())?,
+            self,
+        )
+    }
 }
 
 macro_rules! bin_trait {

candle-core/tests/tensor_tests.rs

@@ -1345,3 +1345,21 @@ fn pow() -> Result<()> {
     );
     Ok(())
 }
+
+#[test]
+fn masked_fill() -> Result<()> {
+    let lhs = Tensor::zeros((5, 5), DType::F32, &Device::Cpu)?;
+    let rhs = Tensor::eye(5, DType::I64, &Device::Cpu)?;
+    let res = lhs.masked_fill(&rhs, f32::NEG_INFINITY)?;
+    assert_eq!(
+        res.to_vec2::<f32>()?,
+        [
+            [f32::NEG_INFINITY, 0.0, 0.0, 0.0, 0.0],
+            [0.0, f32::NEG_INFINITY, 0.0, 0.0, 0.0],
+            [0.0, 0.0, f32::NEG_INFINITY, 0.0, 0.0],
+            [0.0, 0.0, 0.0, f32::NEG_INFINITY, 0.0],
+            [0.0, 0.0, 0.0, 0.0, f32::NEG_INFINITY],
+        ]
+    );
+    Ok(())
+}

candle-transformers/src/models/chatglm.rs

@@ -107,13 +107,6 @@ struct CoreAttention {
     norm_factor: f64,
 }
 
-fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result<Tensor> {
-    let shape = mask.shape();
-    let on_true = Tensor::new(on_true, on_false.device())?.broadcast_as(shape.dims())?;
-    let m = mask.where_cond(&on_true, on_false)?;
-    Ok(m)
-}
-
 impl CoreAttention {
     fn new(layer_number: usize, cfg: &Config) -> Result<Self> {
         let norm_factor = (cfg.kv_channels as f64).sqrt();
@@ -152,8 +145,7 @@ impl CoreAttention {
             Some(coeff) => (matmul_result * coeff)?,
         };
         let attention_scores = match attention_mask {
-            Some(mask) => masked_fill(
-                &matmul_result,
+            Some(mask) => matmul_result.masked_fill(
                 &mask.broadcast_left((matmul_result.dim(0)?, matmul_result.dim(1)?))?,
                 f32::NEG_INFINITY,
             )?,

candle-transformers/src/models/distilbert.rs

@@ -5,13 +5,6 @@ use serde::Deserialize;
 
 pub const DTYPE: DType = DType::F32;
 
-fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result<Tensor> {
-    let shape = mask.shape();
-    let on_true = Tensor::new(on_true, on_false.device())?.broadcast_as(shape.dims())?;
-    let m = mask.where_cond(&on_true, on_false)?;
-    Ok(m)
-}
-
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
 #[serde(rename_all = "lowercase")]
 enum HiddenAct {
@@ -180,7 +173,9 @@ impl MultiHeadSelfAttention {
         let scores = q.matmul(&k.transpose(2, 3)?.contiguous()?)?;
         let mask = attention_mask.broadcast_as(scores.shape())?;
 
-        let scores = masked_fill(&scores.to_dtype(DType::F32)?, &mask, f32::NEG_INFINITY)?;
+        let scores = scores
+            .to_dtype(DType::F32)?
+            .masked_fill(&mask, f32::NEG_INFINITY)?;
         let weights = candle_nn::ops::softmax(&scores, candle::D::Minus1)?;
 
         let context = weights.matmul(&v.contiguous()?)?;

candle-transformers/src/models/falcon.rs

@@ -177,15 +177,6 @@ impl FalconRotaryEmbedding {
     }
 }
 
-fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result<Tensor> {
-    let shape = mask.shape();
-    let on_true = Tensor::new(on_true, on_false.device())?
-        .to_dtype(on_false.dtype())?
-        .broadcast_as(shape.dims())?;
-    let m = mask.where_cond(&on_true, on_false)?;
-    Ok(m)
-}
-
 #[derive(Debug, Clone)]
 struct FalconAttention {
     query_key_value: Linear,
@@ -298,7 +289,9 @@ impl FalconAttention {
         let attention_scores = match mask {
             None => attention_scores,
             Some(mask) => {
-                let mask = masked_fill(&mask.to_dtype(DType::F32)?, mask, -1e9)?
+                let mask = mask
+                    .to_dtype(DType::F32)?
+                    .masked_fill(mask, -1e9)?
                     .to_dtype(query.dtype())?;
                 attention_scores.broadcast_add(&mask.squeeze(1)?)?
             }

candle-transformers/src/models/llama.rs

@@ -255,7 +255,7 @@ impl CausalSelfAttention {
                 att
             } else {
                 let mask = cache.mask(seq_len)?.broadcast_as(att.shape())?;
-                masked_fill(&att, &mask, f32::NEG_INFINITY)?
+                att.masked_fill(&mask, f32::NEG_INFINITY)?
             };
             let att = candle_nn::ops::softmax(&att, D::Minus1)?;
             // Convert to contiguous as matmul doesn't support strided vs for now.
@@ -295,13 +295,6 @@ impl CausalSelfAttention {
     }
 }
 
-fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result<Tensor> {
-    let shape = mask.shape();
-    let on_true = Tensor::new(on_true, on_false.device())?.broadcast_as(shape.dims())?;
-    let m = mask.where_cond(&on_true, on_false)?;
-    Ok(m)
-}
-
 #[derive(Debug, Clone)]
 struct Mlp {
     c_fc1: Linear,

candle-transformers/src/models/llama2_c.rs

@@ -198,7 +198,7 @@ impl CausalSelfAttention {
             att
         } else {
             let mask = cache.mask(seq_len)?.broadcast_as(att.shape())?;
-            masked_fill(&att, &mask, f32::NEG_INFINITY)?
+            att.masked_fill(&mask, f32::NEG_INFINITY)?
         };
         let att = candle_nn::ops::softmax(&att, D::Minus1)?;
         // Convert to contiguous as matmul doesn't support strided vs for now.
@@ -242,13 +242,6 @@ impl CausalSelfAttention {
     }
 }
 
-fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result<Tensor> {
-    let shape = mask.shape();
-    let on_true = Tensor::new(on_true, on_false.device())?.broadcast_as(shape.dims())?;
-    let m = mask.where_cond(&on_true, on_false)?;
-    Ok(m)
-}
-
 #[derive(Debug, Clone)]
 struct Mlp {
     c_fc1: Linear,

candle-transformers/src/models/mpt.rs

@@ -118,11 +118,8 @@ impl GroupedQueryAttention {
         let attn_weights = attn_weights.broadcast_add(&attn_bias)?;
         let attn_weights = match mask {
             None => attn_weights,
-            Some(mask) => masked_fill(
-                &attn_weights,
-                &mask.broadcast_as(attn_weights.shape())?,
-                f32::NEG_INFINITY,
-            )?,
+            Some(mask) => attn_weights
+                .masked_fill(&mask.broadcast_as(attn_weights.shape())?, f32::NEG_INFINITY)?,
         };
         let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
         let attn_output = attn_weights
@@ -281,10 +278,3 @@ pub(crate) fn get_mask(size: usize, device: &Device) -> Result<Tensor> {
         .collect();
     Tensor::from_slice(&mask, (size, size), device)
 }
-
-pub(crate) fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result<Tensor> {
-    let shape = mask.shape();
-    let on_true = Tensor::new(on_true, on_false.device())?.broadcast_as(shape.dims())?;
-    let m = mask.where_cond(&on_true, on_false)?;
-    Ok(m)
-}

candle-transformers/src/models/phi.rs

@@ -126,13 +126,6 @@ fn get_mask(size: usize, device: &Device) -> Result<Tensor> {
     Tensor::from_slice(&mask, (size, size), device)
 }
 
-fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result<Tensor> {
-    let shape = mask.shape();
-    let on_true = Tensor::new(on_true, on_false.device())?.broadcast_as(shape.dims())?;
-    let m = mask.where_cond(&on_true, on_false)?;
-    Ok(m)
-}
-
 impl Attention {
     fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
         let num_heads = cfg.num_attention_heads;
@@ -233,8 +226,7 @@ impl Attention {
             * self.softmax_scale)?;
         let attn_weights = match mask {
             None => attn_weights,
-            Some(mask) => masked_fill(
-                &attn_weights,
+            Some(mask) => attn_weights.masked_fill(
                 &mask.broadcast_left((b_size, self.num_heads))?,
                 f32::NEG_INFINITY,
             )?,

candle-transformers/src/models/quantized_llama.rs

@@ -138,19 +138,12 @@ struct LayerWeights {
     head_dim: usize,
     cos: Tensor,
     sin: Tensor,
-    neg_inf: Tensor,
     kv_cache: Option<(Tensor, Tensor)>,
     span_attn: tracing::Span,
     span_rot: tracing::Span,
     span_mlp: tracing::Span,
 }
 
-fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: &Tensor) -> Result<Tensor> {
-    let shape = mask.shape();
-    let m = mask.where_cond(&on_true.broadcast_as(shape.dims())?, on_false)?;
-    Ok(m)
-}
-
 impl LayerWeights {
     fn apply_rotary_emb(&self, x: &Tensor, index_pos: usize) -> Result<Tensor> {
         let _enter = self.span_rot.enter();
@@ -214,7 +207,7 @@ impl LayerWeights {
             None => att,
             Some(mask) => {
                 let mask = mask.broadcast_as(att.shape())?;
-                masked_fill(&att, &mask, &self.neg_inf)?
+                att.masked_fill(&mask, f32::NEG_INFINITY)?
             }
         };
         let att = candle_nn::ops::softmax_last_dim(&att)?;
@@ -260,7 +253,6 @@ impl ModelWeights {
     pub fn from_ggml(mut ct: ggml_file::Content, gqa: usize) -> Result<Self> {
         let head_dim = (ct.hparams.n_embd / ct.hparams.n_head) as usize;
         let (cos, sin) = precomput_freqs_cis(head_dim, 10000., &ct.device)?;
-        let neg_inf = Tensor::new(f32::NEG_INFINITY, &ct.device)?;
         let tok_embeddings = ct.remove("tok_embeddings.weight")?;
         let tok_embeddings = tok_embeddings.dequantize(&ct.device)?;
         let norm = RmsNorm::from_qtensor(ct.remove("norm.weight")?, 1e-5)?;
@@ -300,7 +292,6 @@ impl ModelWeights {
                 head_dim: (ct.hparams.n_embd / ct.hparams.n_head) as usize,
                 cos: cos.clone(),
                 sin: sin.clone(),
-                neg_inf: neg_inf.clone(),
                 kv_cache: None,
                 span_attn,
                 span_rot,
@@ -349,7 +340,6 @@ impl ModelWeights {
             .and_then(|m| m.to_f32())
             .unwrap_or(10000f32);
         let (cos, sin) = precomput_freqs_cis(rope_dim, rope_freq_base, device)?;
-        let neg_inf = Tensor::new(f32::NEG_INFINITY, device)?;
 
         let tok_embeddings = ct.tensor(reader, "token_embd.weight", device)?;
         let tok_embeddings = tok_embeddings.dequantize(device)?;
@@ -420,7 +410,6 @@ impl ModelWeights {
                 head_dim: embedding_length / head_count,
                 cos: cos.clone(),
                 sin: sin.clone(),
-                neg_inf: neg_inf.clone(),
                 kv_cache: None,
                 span_attn,
                 span_rot,

candle-transformers/src/models/quantized_llama2_c.rs

@@ -75,7 +75,7 @@ impl CausalSelfAttention {
             att
         } else {
             let mask = cache.mask(seq_len)?.broadcast_as(att.shape())?;
-            masked_fill(&att, &mask, f32::NEG_INFINITY)?
+            att.masked_fill(&mask, f32::NEG_INFINITY)?
         };
         let att = candle_nn::ops::softmax(&att, D::Minus1)?;
         // Convert to contiguous as matmul doesn't support strided vs for now.
@@ -119,13 +119,6 @@ impl CausalSelfAttention {
     }
 }
 
-fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result<Tensor> {
-    let shape = mask.shape();
-    let on_true = Tensor::new(on_true, on_false.device())?.broadcast_as(shape.dims())?;
-    let m = mask.where_cond(&on_true, on_false)?;
-    Ok(m)
-}
-
 #[derive(Debug, Clone)]
 struct Mlp {
     c_fc1: Linear,

candle-transformers/src/models/quantized_mixformer.rs

@@ -32,13 +32,6 @@ fn get_mask(size: usize, device: &Device) -> Result<Tensor> {
     Tensor::from_slice(&mask, (size, size), device)
 }
 
-fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: f32) -> Result<Tensor> {
-    let shape = mask.shape();
-    let on_true = Tensor::new(on_true, on_false.device())?.broadcast_as(shape.dims())?;
-    let m = mask.where_cond(&on_true, on_false)?;
-    Ok(m)
-}
-
 #[derive(Debug, Clone)]
 struct RotaryEmbedding {
     sin: Tensor,
@@ -219,8 +212,7 @@ impl MHA {
         // scores = scores + causal_mask.to(dtype=scores.dtype)
         let attn_weights = match mask {
             None => attn_weights,
-            Some(mask) => masked_fill(
-                &attn_weights,
+            Some(mask) => attn_weights.masked_fill(
                 &mask.broadcast_left(b_size * self.n_head)?,
                 f32::NEG_INFINITY,
             )?,

candle-transformers/src/models/quantized_mpt.rs

@@ -85,11 +85,8 @@ impl GroupedQueryAttention {
         let attn_weights = attn_weights.broadcast_add(&attn_bias)?;
         let attn_weights = match mask {
             None => attn_weights,
-            Some(mask) => super::mpt::masked_fill(
-                &attn_weights,
-                &mask.broadcast_as(attn_weights.shape())?,
-                f32::NEG_INFINITY,
-            )?,
+            Some(mask) => attn_weights
+                .masked_fill(&mask.broadcast_as(attn_weights.shape())?, f32::NEG_INFINITY)?,
         };
         let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
         let attn_output = attn_weights

candle-transformers/src/models/quantized_phi.rs

@@ -61,18 +61,11 @@ struct LayerWeights {
     cos: Tensor,
     sin: Tensor,
     rope_dim: usize,
-    neg_inf: Tensor,
     kv_cache: Option<(Tensor, Tensor)>,
     span_attn: tracing::Span,
     span_rot: tracing::Span,
 }
 
-fn masked_fill(on_false: &Tensor, mask: &Tensor, on_true: &Tensor) -> Result<Tensor> {
-    let shape = mask.shape();
-    let m = mask.where_cond(&on_true.broadcast_as(shape.dims())?, on_false)?;
-    Ok(m)
-}
-
 impl LayerWeights {
     fn apply_rotary_emb(&self, xs: &Tensor, index_pos: usize) -> Result<Tensor> {
         let _enter = self.span_rot.enter();
@@ -131,7 +124,7 @@ impl LayerWeights {
             None => att,
             Some(mask) => {
                 let mask = mask.broadcast_as(att.shape())?;
-                masked_fill(&att, &mask, &self.neg_inf)?
+                att.masked_fill(&mask, f32::NEG_INFINITY)?
             }
         };
         let att = candle_nn::ops::softmax_last_dim(&att)?;
@@ -199,7 +192,6 @@ impl ModelWeights {
         let rope_dim = md_get("phi2.rope.dimension_count")?.to_u32()? as usize;
         let ln_eps = md_get("phi2.attention.layer_norm_epsilon")?.to_f32()? as f64;
         let (cos, sin) = precomput_freqs_cis(rope_dim, 10_000., device)?;
-        let neg_inf = Tensor::new(f32::NEG_INFINITY, device)?;
 
         let tok_embeddings = ct.tensor(reader, "token_embd.weight", device)?;
         let tok_embeddings = tok_embeddings.dequantize(device)?;
@@ -233,7 +225,6 @@ impl ModelWeights {
                 cos: cos.clone(),
                 sin: sin.clone(),
                 rope_dim,
-                neg_inf: neg_inf.clone(),
                 kv_cache: None,
                 span_attn,
                 span_rot,