mistr_mistral_block — mistr_mistral

A step of the mistr_native implementation. No other function in the package calls it. See the file header for the source the module follows.

Usage

mistr_mistral_block(
  X,
  Wq,
  Wk,
  Wv,
  Wo,
  W1,
  W2,
  W3,
  n_heads,
  n_kv_heads,
  window,
  norm1 = NULL,
  norm2 = NULL,
  base = 10000
)

Arguments

X: A matrix; passed to as.matrix.
Wq: Passed to proj.
Wk: Passed to proj.
Wv: Passed to proj.
Wo: Passed to proj.
W1: Passed to mistr_swiglu.
W2: Passed to mistr_swiglu.
W3: Passed to mistr_swiglu.
n_heads: Carried through into a list the body builds.
n_kv_heads: Numeric; combined arithmetically in the body.
window: Coerced to integer by the body, with as.integer.
norm1: Optional; may be NULL. Passed to mistr_rms_norm.
norm2: Passed to mistr_rms_norm.
base: Passed to mistr_grouped_query_attention. Defaults to 10000.

Value

A list with estimate, output, attention_mask, L, d, n_heads, n_kv_heads, window, kv_cache_entries, method.

Examples

set.seed(15)
L <- 4; d <- 4; dff <- 6
X <- matrix(rnorm(L * d), L, d)
mk <- function(r, c) matrix(rnorm(r * c, sd = 0.3), r, c)
b <- mistr_mistral_block(X, Wq = mk(d, d), Wk = mk(d, d / 2), Wv = mk(d, d / 2), Wo = mk(d, d),
                         W1 = mk(d, dff), W2 = mk(dff, d), W3 = mk(d, dff),
                         n_heads = 2, n_kv_heads = 1, window = 2)
dim(b$output)
#> [1] 4 4