Struct QueryKeyRotaryEmbeddings

pub struct QueryKeyRotaryEmbeddings {
    pub head_width: usize,
    /* private fields */
}

Fields

head_width: usize

Implementations

impl QueryKeyRotaryEmbeddings

pub fn forward( &self, query: &Tensor, key: &Tensor, cache: &LayerKeyValueCache, positions: Option<&Tensor>, ) -> Result<(Tensor, Tensor), QueryKeyRotaryEmbeddingsError>

Construct query-key rotary embeddings module.

The rotary embedding will be precomputed for up to seq_len positions. The embedding will be recomputed when a longer sequence is found in the input.

• fraction - Fraction of hidden width to apply rotary embeddings to. Must be in [0,1].

• head_width - Query and key attention head width.

• seq_len - Number of positions to initially precompute.

• base - The base used for theta (normally 10_000). Determines the cycle length of the embeddings.

• device - Device on which the module is to be initialized. Apply rotary embeddings to the query and key.

• query - Query to apply the rotary embeddings to. Shape: (batch_size, n_heads, seq_len, width_per_head)

• key - Key to apply the rotary embeddings to. Shape: (batch_size, n_heads, seq_len, width_per_head)

• positions - Positions of the inputs. If no positions are provided, they are assumed to be [0, seq_len). Shape: (batch_size, seq_len)

Returns: Input with the rotary embeddings applied. Shape: (batch_size, n_heads, seq_len, width_per_head)

Trait Implementations

impl Debug for QueryKeyRotaryEmbeddings

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.