# TRITIUM-PROVENANCE sha256=6d7ce1815a5ea4eacf5df829ac23026deec566513c6bfb4d5c0fbfd4974218ee lang=python
# TRITIUM-HEIGHT high=0.19 mid=0.29 low=0.52 fns=17
# TRITIUM-PSEUDO fn=zip_add
#| fn zip_add(a, b, n, out):
#|   return n
# TRITIUM-HEIGHT fn=zip_add high=0.00 mid=0.23 low=0.77 role=mid_helper
def zip_add(a, b, n, out):
    i = 0
    while i < n:
        out[i] = a[i] + b[i]
        i = i + 1
    return n

# TRITIUM-PSEUDO fn=sum_f
#| fn sum_f(vals, n):
#|   return s
# TRITIUM-HEIGHT fn=sum_f high=0.00 mid=0.41 low=0.59 role=mid_helper
def sum_f(vals, n):
    s = 0.0
    i = 0
    while i < n:
        s = s + vals[i]
        i = i + 1
    return s

# TRITIUM-PSEUDO fn=fmean
#| fn fmean(xs, n):
#|   return sum_f(xs, n) / float(n)
# TRITIUM-HEIGHT fn=fmean high=0.75 mid=0.00 low=0.25
def fmean(xs: list[float], n: int) -> float:
    return sum_f(xs, n) / float(n)

# TRITIUM-PSEUDO fn=fsqrt
#| fn fsqrt(x):
#|   return g
# TRITIUM-HEIGHT fn=fsqrt high=0.00 mid=0.63 low=0.37 role=algo_mid
def fsqrt(x: float) -> float:
    if x <= 0.0:
        return 0.0
    g = x
    i = 0
    while i < 8:
        g = 0.5 * (g + x / g)
        i = i + 1
    return g

# TRITIUM-PSEUDO fn=layer_norm
#| fn layer_norm(xs, n, out, eps):
#|   fmean(...)
#|   float(...)
#|   fsqrt(...)
#|   return mu
# TRITIUM-HEIGHT fn=layer_norm high=0.24 mid=0.29 low=0.47 role=algo_mid
def layer_norm(xs: list[float], n: int, out: list[float], eps: float) -> float:
    mu = fmean(xs, n)
    v = 0.0
    i = 0
    while i < n:
        d = xs[i] - mu
        v = v + d * d
        i = i + 1
    v = v / float(n)
    inv = 1.0 / fsqrt(v + eps)
    i = 0
    while i < n:
        out[i] = (xs[i] - mu) * inv
        i = i + 1
    return mu

# TRITIUM-PSEUDO fn=ln_rows
#| fn ln_rows(x, n_tokens, d, out, tmp, nrm, eps):
#|   layer_norm(...)
#|   float(...)
#|   return float(n_tokens) * 1.0
# TRITIUM-HEIGHT fn=ln_rows high=0.14 mid=0.19 low=0.66 role=algo_mid
def ln_rows(x: list[float], n_tokens: int, d: int, out: list[float], tmp: list[float], nrm: list[float], eps: float) -> float:
    t = 0
    while t < n_tokens:
        i = 0
        while i < d:
            tmp[i] = x[t * d + i]
            i = i + 1
        layer_norm(tmp, d, nrm, eps)
        i = 0
        while i < d:
            out[t * d + i] = nrm[i]
            i = i + 1
        t = t + 1
    return float(n_tokens) * 1.0

# TRITIUM-PSEUDO fn=gemv_f
#| fn gemv_f(a, x, n, m, out):
#|   return n
# TRITIUM-HEIGHT fn=gemv_f high=0.00 mid=0.26 low=0.74 role=algo_mid
def gemv_f(a: list[float], x: list[float], n: int, m: int, out: list[float]) -> int:
    i = 0
    while i < n:
        s = 0.0
        j = 0
        while j < m:
            s = s + a[i * m + j] * x[j]
            j = j + 1
        out[i] = s
        i = i + 1
    return n

# TRITIUM-PSEUDO fn=qkv_seq
#| fn qkv_seq(x, n_tokens, d, w, out, tmp, row):
#|   gemv_f(...)
#|   return n_tokens
# TRITIUM-HEIGHT fn=qkv_seq high=0.07 mid=0.20 low=0.73 role=algo_mid
def qkv_seq(x: list[float], n_tokens: int, d: int, w: list[float], out: list[float], tmp: list[float], row: list[float]) -> int:
    t = 0
    while t < n_tokens:
        i = 0
        while i < d:
            tmp[i] = x[t * d + i]
            i = i + 1
        gemv_f(w, tmp, d, d, row)
        i = 0
        while i < d:
            out[t * d + i] = row[i]
            i = i + 1
        t = t + 1
    return n_tokens

# TRITIUM-PSEUDO fn=pack_heads
#| fn pack_heads(src, n_tokens, heads, dh, dest):
#|   return heads * n_tokens * dh
# TRITIUM-HEIGHT fn=pack_heads high=0.00 mid=0.27 low=0.73 role=algo_mid
def pack_heads(src: list[float], n_tokens: int, heads: int, dh: int, dest: list[float]) -> int:
    h = 0
    while h < heads:
        t = 0
        while t < n_tokens:
            u = 0
            while u < dh:
                dest[h * n_tokens * dh + t * dh + u] = src[t * heads * dh + h * dh + u]
                u = u + 1
            t = t + 1
        h = h + 1
    return heads * n_tokens * dh

# TRITIUM-PSEUDO fn=fexp
#| fn fexp(x):
#|   float(...)
# TRITIUM-HEIGHT fn=fexp high=0.12 mid=0.66 low=0.22 role=algo_mid
def fexp(x: float) -> float:
    ln2 = 0.6931471805599453
    inv_ln2 = 1.4426950408889634
    if x >= 0.0:
        choice = 0.5
    else:
        choice = -0.5
    n = int(x * inv_ln2 + choice)
    if n > 60:
        n = 60
    if n < -60:
        n = -60
    r = x - float(n) * ln2
    s = 1.0
    t = 1.0
    i = 1
    while i <= 14:
        t = t * r / float(i)
        s = s + t
        i = i + 1
    p = 1.0
    nn = n
    if nn < 0:
        nn = 0 - nn
        j = 0
        while j < nn:
            p = p * 2.0
            j = j + 1
        return s / p
    else:
        j = 0
        while j < nn:
            p = p * 2.0
            j = j + 1
        return s * p

# TRITIUM-PSEUDO fn=mha_fwd
#| fn mha_fwd(q, k, v, scores, out, h, l, d, scale):
#|   fexp(...)
#|   return h * l * d
# TRITIUM-HEIGHT fn=mha_fwd high=0.03 mid=0.32 low=0.66 role=algo_mid
def mha_fwd(q: list[float], k: list[float], v: list[float], scores: list[float], out: list[float], h: int, l: int, d: int, scale: float) -> int:
    heap_dist = 0
    while heap_dist < h:
        i = 0
        while i < l:
            j = 0
            while j < l:
                s = 0.0
                t = 0
                while t < d:
                    s = s + q[heap_dist * l * d + i * d + t] * k[heap_dist * l * d + j * d + t]
                    t = t + 1
                scores[i * l + j] = s * scale
                j = j + 1
            i = i + 1
        i = 0
        while i < l:
            m = scores[i * l]
            j = 1
            while j < l:
                if scores[i * l + j] > m:
                    m = scores[i * l + j]
                j = j + 1
            ssum = 0.0
            j = 0
            while j < l:
                e = fexp(scores[i * l + j] - m)
                scores[i * l + j] = e
                ssum = ssum + e
                j = j + 1
            j = 0
            while j < l:
                scores[i * l + j] = scores[i * l + j] / ssum
                j = j + 1
            i = i + 1
        i = 0
        while i < l:
            t = 0
            while t < d:
                s = 0.0
                j = 0
                while j < l:
                    s = s + scores[i * l + j] * v[heap_dist * l * d + j * d + t]
                    j = j + 1
                out[heap_dist * l * d + i * d + t] = s
                t = t + 1
            i = i + 1
        heap_dist = heap_dist + 1
    return h * l * d

# TRITIUM-PSEUDO fn=unpack_heads
#| fn unpack_heads(src, n_tokens, heads, dh, dest):
#|   return n_tokens * heads * dh
# TRITIUM-HEIGHT fn=unpack_heads high=0.00 mid=0.27 low=0.73 role=algo_mid
def unpack_heads(src: list[float], n_tokens: int, heads: int, dh: int, dest: list[float]) -> int:
    h = 0
    while h < heads:
        t = 0
        while t < n_tokens:
            u = 0
            while u < dh:
                dest[t * heads * dh + h * dh + u] = src[h * n_tokens * dh + t * dh + u]
                u = u + 1
            t = t + 1
        h = h + 1
    return n_tokens * heads * dh

# TRITIUM-PSEUDO fn=residual_add
#| fn residual_add(a, b, n, out):
#|   return zip_add(a, b, n, out)
# TRITIUM-HEIGHT fn=residual_add high=0.33 mid=0.00 low=0.67
def residual_add(a: list[float], b: list[float], n: int, out: list[float]) -> int:
    return zip_add(a, b, n, out)

# TRITIUM-PSEUDO fn=ftanh
#| fn ftanh(x):
#|   fexp(...)
#|   return (e - 1.0) / (e + 1.0)
# TRITIUM-HEIGHT fn=ftanh high=0.25 mid=0.75 low=0.00
def ftanh(x: float) -> float:
    a = x
    if a > 20.0:
        a = 20.0
    if a < -20.0:
        a = -20.0
    e = fexp(a + a)
    return (e - 1.0) / (e + 1.0)

# TRITIUM-PSEUDO fn=gelu_tanh
#| fn gelu_tanh(x):
#|   x3 = x * x * x
#|   return 0.5 * x * (1.0 + ftanh(0.79788456 * (x + 0.044715 * x3)))
# TRITIUM-HEIGHT fn=gelu_tanh high=0.75 mid=0.25 low=0.00
def gelu_tanh(x: float) -> float:
    x3 = x * x * x
    return 0.5 * x * (1.0 + ftanh(0.79788456 * (x + 0.044715 * x3)))

# TRITIUM-PSEUDO fn=mlp_gelu_vec
#| fn mlp_gelu_vec(x, w1, w2, d, hidden, hidden_out, out):
#|   gemv_f(...)
#|   gelu_tanh(...)
#|   return d
# TRITIUM-HEIGHT fn=mlp_gelu_vec high=0.30 mid=0.10 low=0.60
def mlp_gelu_vec(x: list[float], w1: list[float], w2: list[float], d: int, hidden: int, hidden_out: list[float], out: list[float]) -> int:
    gemv_f(w1, x, hidden, d, hidden_out)
    i = 0
    while i < hidden:
        hidden_out[i] = gelu_tanh(hidden_out[i])
        i = i + 1
    gemv_f(w2, hidden_out, d, hidden, out)
    return d

# TRITIUM-PSEUDO fn=transformer_block
#| fn transformer_block(x, n_tokens, d, heads, dh, hidden, scale, wq, wk, wv, wo, w1, w2, q, k, v, qh, kh, vh, scores, attention, flat, normed, normed2, tmp, nrm, row, hidden_out, mlp_out, projected, residual, out):
#|   ln_rows(...)
#|   qkv_seq(...)
#|   pack_heads(...)
#|   mha_fwd(...)
#|   unpack_heads(...)
#|   residual_add(...)
#|   gemv_f(...)
#|   mlp_gelu_vec(...)
#|   return n_tokens
# TRITIUM-HEIGHT fn=transformer_block high=0.28 mid=0.11 low=0.61
def transformer_block(x: list[float], n_tokens: int, d: int, heads: int, dh: int, hidden: int, scale: float, wq: list[float], wk: list[float], wv: list[float], wo: list[float], w1: list[float], w2: list[float], q: list[float], k: list[float], v: list[float], qh: list[float], kh: list[float], vh: list[float], scores: list[float], attention: list[float], flat: list[float], normed: list[float], normed2: list[float], tmp: list[float], nrm: list[float], row: list[float], hidden_out: list[float], mlp_out: list[float], projected: list[float], residual: list[float], out: list[float]) -> int:
    ln_rows(x, n_tokens, d, normed, tmp, nrm, 0.0)
    qkv_seq(normed, n_tokens, d, wq, q, tmp, row)
    qkv_seq(normed, n_tokens, d, wk, k, tmp, row)
    qkv_seq(normed, n_tokens, d, wv, v, tmp, row)
    pack_heads(q, n_tokens, heads, dh, qh)
    pack_heads(k, n_tokens, heads, dh, kh)
    pack_heads(v, n_tokens, heads, dh, vh)
    mha_fwd(qh, kh, vh, scores, attention, heads, n_tokens, dh, scale)
    unpack_heads(attention, n_tokens, heads, dh, flat)
    t = 0
    while t < n_tokens:
        i = 0
        while i < d:
            tmp[i] = flat[t * d + i]
            i = i + 1
        gemv_f(wo, tmp, d, d, row)
        i = 0
        while i < d:
            projected[t * d + i] = row[i]
            i = i + 1
        t = t + 1
    residual_add(x, projected, n_tokens * d, residual)
    ln_rows(residual, n_tokens, d, normed2, tmp, nrm, 0.0)
    t = 0
    while t < n_tokens:
        i = 0
        while i < d:
            tmp[i] = normed2[t * d + i]
            i = i + 1
        mlp_gelu_vec(tmp, w1, w2, d, hidden, hidden_out, row)
        i = 0
        while i < d:
            mlp_out[t * d + i] = row[i]
            i = i + 1
        t = t + 1
    residual_add(residual, mlp_out, n_tokens * d, out)
    return n_tokens