Struct CausalMasker

pub struct CausalMasker;

Implementations

impl CausalMasker

pub fn expand_mask( &self, mask: &Tensor, dtype: DType, tgt_len: Option<usize>, ) -> Result<Tensor, Error>

Expands a mask from (bs, seq_len) to (bs, 1, tgt_len, seq_len) If tgt_len is None, use seq_len

pub fn calculate_past_kv_len( &self, cache: &[Option<(Tensor, Tensor)>], ) -> Result<usize, Error>

pub fn make_causal_mask_matrix( &self, input_ids: &Tensor, cache: &dyn PastKvLenCache, dtype: DType, n_attn_heads: usize, ) -> Result<Option<Tensor>, Error>

pub fn make_sliding_window_causal_mask_matrix( &self, input_ids: &Tensor, cache: &dyn PastKvLenCache, sliding_window: Option<usize>, dtype: DType, n_attn_heads: usize, ) -> Result<Option<Tensor>, Error>

pub fn make_causal_mask_as_attn_bias( &self, input_ids: &Tensor, cache: &dyn PastKvLenCache, dtype: DType, n_attn_heads: usize, ) -> Result<Option<Tensor>, Error>

👎Deprecated since 0.3.4: use make_causal_mask_matrix_as_attn_bias instead. This is incompatible with Sdpa.

pub fn make_causal_mask_with_sliding_window_as_attn_bias( &self, input_ids: &Tensor, cache: &dyn PastKvLenCache, sliding_window: Option<usize>, dtype: DType, n_attn_heads: usize, ) -> Result<Option<Tensor>, Error>

👎Deprecated since 0.3.4: use make_causal_mask_matrix_with_sliding_window_as_attn_bias instead. This is incompatible with Sdpa.

pub fn apply_mask_one_and_zero( &self, mask: &Option<Tensor>, att: Tensor, neg_inf: &Tensor, ) -> Result<Tensor, Error>