/* TRITIUM-PROVENANCE sha256=6d7ce1815a5ea4eacf5df829ac23026deec566513c6bfb4d5c0fbfd4974218ee lang=c */
/* TRITIUM-HEIGHT high=0.13 mid=0.43 low=0.44 fns=15 */
#include <stdint.h>
#include <math.h>

// TRITIUM-PSEUDO fn=fmean
//| fn fmean(vals, n):
//|   return x / (double)(n)
// TRITIUM-HEIGHT fn=fmean high=0.00 mid=0.41 low=0.59
double fmean(double* vals, int64_t n) {
    int64_t i;
    double x;

    x = 0.0;
    i = 0L;
    while (i < n) {
        x = x + vals[i];
        i = (i + 1L);
    }
    return x / (double)(n);
}

// TRITIUM-PSEUDO fn=fsqrt
//| fn fsqrt(x):
//|     return 0.0
//|     return i_2
// TRITIUM-HEIGHT fn=fsqrt high=0.00 mid=0.69 low=0.31
double fsqrt(double x) {
    int64_t i;
    int fcmp4;
    double i_2;
    double x_2;

    fcmp4 = x <= 0.0;
    if (fcmp4) {
        return 0.0;
    } else {
        i_2 = x;
        i = 0L;
        while (i < 8L) {
            x_2 = 0.5 * (i_2 + (x / i_2));
            i_2 = x_2;
            i = (i + 1L);
        }
        return i_2;
    }
}

// TRITIUM-PSEUDO fn=layer_norm
//| fn layer_norm(vals, n, out, eps):
//|   fmean(...)
//|   fsqrt(...)
//|   return mu
// TRITIUM-HEIGHT fn=layer_norm high=0.20 mid=0.36 low=0.44
double layer_norm(double* vals, int64_t n, double* out, double eps) {
    int64_t i;
    double mu;
    double x;
    double x_2;
    double x_3;

    mu = fmean(vals, n);
    x_2 = 0.0;
    i = 0L;
    while (i < n) {
        x = vals[i] - mu;
        x_2 = x_2 + (x * x);
        i = (i + 1L);
    }
    x_2 = (x_2 / (double)(n));
    x_3 = 1.0 / (fsqrt(x_2 + eps));
    i = 0L;
    while (i < n) {
        out[i] = ((vals[i] - mu) * x_3);
        i = (i + 1L);
    }
    return mu;
}

// TRITIUM-PSEUDO fn=ln_rows
//| fn ln_rows(x, n_tokens, d, out, tmp, nrm, eps):
//|   layer_norm(...)
//|   return x_2
// TRITIUM-HEIGHT fn=ln_rows high=0.11 mid=0.33 low=0.56
double ln_rows(double* x, int64_t n_tokens, int64_t d, double* out, double* tmp, double* nrm, double eps) {
    int64_t t;
    int64_t i;
    double x_2;

    t = 0L;
    while (t < n_tokens) {
        i = 0L;
        while (i < d) {
            tmp[i] = x[((t * d) + i)];
            i = (i + 1L);
        }
        layer_norm(tmp, d, nrm, eps);
        i = 0L;
        while (i < d) {
            out[((t * d) + i)] = nrm[i];
            i = (i + 1L);
        }
        t = (t + 1L);
    }
    x_2 = (double)(n_tokens) * 1.0;
    return x_2;
}

// 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.38 low=0.62
int64_t gemv_f(double* a, double* x, int64_t n, int64_t m, double* out) {
    int64_t i;
    int64_t j;
    double x_2;

    i = 0L;
    while (i < n) {
        x_2 = 0.0;
        j = 0L;
        while (j < m) {
            x_2 = x_2 + (a[((i * m) + j)] * x[j]);
            j = (j + 1L);
        }
        out[i] = x_2;
        i = (i + 1L);
    }
    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.11 mid=0.30 low=0.58
int64_t qkv_seq(double* x, int64_t n_tokens, int64_t d, double* w, double* out, double* tmp, double* row) {
    int64_t t;
    int64_t i;

    t = 0L;
    while (t < n_tokens) {
        i = 0L;
        while (i < d) {
            tmp[i] = x[((t * d) + i)];
            i = (i + 1L);
        }
        gemv_f(w, tmp, d, d, row);
        i = 0L;
        while (i < d) {
            out[((t * d) + i)] = row[i];
            i = (i + 1L);
        }
        t = (t + 1L);
    }
    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.38 low=0.62
int64_t pack_heads(double* src, int64_t n_tokens, int64_t heads, int64_t dh, double* dest) {
    int64_t h;
    int64_t t;
    int64_t u;

    h = 0L;
    while (h < heads) {
        t = 0L;
        while (t < n_tokens) {
            u = 0L;
            while (u < dh) {
                dest[((((h * n_tokens) * dh) + (t * dh)) + u)] = src[((((t * heads) * dh) + (h * dh)) + u)];
                u = (u + 1L);
            }
            t = (t + 1L);
        }
        h = (h + 1L);
    }
    return (heads * n_tokens) * dh;
}

// TRITIUM-PSEUDO fn=fexp
//| fn fexp(x):
//|     return x_5 / x_4
//|     return x_5 * x_4
// TRITIUM-HEIGHT fn=fexp high=0.00 mid=0.76 low=0.24
double fexp(double x) {
    double ln2;
    double inv_ln2;
    int64_t i;
    int64_t j;
    int fcmp15;
    int64_t ifexp_2;
    double ifexp;
    int64_t xi;
    double x_2;
    double x_5;
    int64_t i_2;
    double x_4;
    double x_3;

    ln2 = 0.6931471805599453;
    inv_ln2 = 1.4426950408889634;
    fcmp15 = x >= 0.0;
    if (fcmp15) {
        ifexp_2 = 0.5;
    } else {
        ifexp_2 = -0.5;
    }
    ifexp = ifexp_2;
    xi = (int64_t)(((x * inv_ln2) + ifexp));
    if (xi > 60L) {
        xi = 60L;
    }
    if (xi < -60L) {
        xi = -60L;
    }
    x_5 = 1.0;
    x_2 = 1.0;
    i = 1L;
    while (i <= 14L) {
        x_2 = (x_2 * (x - ((double)(xi) * ln2))) / (double)(i);
        x_5 = x_5 + x_2;
        i = (i + 1L);
    }
    x_4 = 1.0;
    i_2 = xi;
    if (i_2 < 0L) {
        i_2 = (0L - i_2);
        j = 0L;
        while (j < i_2) {
            x_4 = x_4 * 2.0;
            j = (j + 1L);
        }
        return x_5 / x_4;
    } else {
        j = 0L;
        while (j < i_2) {
            x_3 = x_4 * 2.0;
            x_4 = x_3;
            j = (j + 1L);
        }
        return x_5 * x_4;
    }
}

// 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.04 mid=0.47 low=0.50
int64_t mha_fwd(double* q, double* k, double* v, double* scores, double* out, int64_t h, int64_t l, int64_t d, double scale) {
    int64_t heap_dist;
    int64_t i;
    int64_t j;
    int64_t t;
    double x_2;
    double i_2;
    int fcmp96;
    double e;
    double x;

    heap_dist = 0L;
    while (heap_dist < h) {
        i = 0L;
        while (i < l) {
            j = 0L;
            while (j < l) {
                x_2 = 0.0;
                t = 0L;
                while (t < d) {
                    x_2 = x_2 + (q[((((heap_dist * l) * d) + (i * d)) + t)] * k[((((heap_dist * l) * d) + (j * d)) + t)]);
                    t = (t + 1L);
                }
                scores[((i * l) + j)] = (x_2 * scale);
                j = (j + 1L);
            }
            i = (i + 1L);
        }
        i = 0L;
        while (i < l) {
            i_2 = scores[(i * l)];
            j = 1L;
            while (j < l) {
                fcmp96 = scores[((i * l) + j)] > i_2;
                if (fcmp96) {
                    i_2 = scores[((i * l) + j)];
                }
                j = (j + 1L);
            }
            x = 0.0;
            j = 0L;
            while (j < l) {
                e = fexp((scores[((i * l) + j)] - i_2));
                scores[((i * l) + j)] = e;
                x = x + e;
                j = (j + 1L);
            }
            j = 0L;
            while (j < l) {
                scores[((i * l) + j)] = (scores[((i * l) + j)] / x);
                j = (j + 1L);
            }
            i = (i + 1L);
        }
        i = 0L;
        while (i < l) {
            t = 0L;
            while (t < d) {
                x_2 = 0.0;
                j = 0L;
                while (j < l) {
                    x_2 = (x_2 + (scores[((i * l) + j)] * v[((((heap_dist * l) * d) + (j * d)) + t)]));
                    j = (j + 1L);
                }
                out[((((heap_dist * l) * d) + (i * d)) + t)] = x_2;
                t = (t + 1L);
            }
            i = (i + 1L);
        }
        heap_dist = (heap_dist + 1L);
    }
    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.38 low=0.62
int64_t unpack_heads(double* src, int64_t n_tokens, int64_t heads, int64_t dh, double* dest) {
    int64_t h;
    int64_t t;
    int64_t u;

    h = 0L;
    while (h < heads) {
        t = 0L;
        while (t < n_tokens) {
            u = 0L;
            while (u < dh) {
                dest[((((t * heads) * dh) + (h * dh)) + u)] = src[((((h * n_tokens) * dh) + (t * dh)) + u)];
                u = (u + 1L);
            }
            t = (t + 1L);
        }
        h = (h + 1L);
    }
    return (n_tokens * heads) * dh;
}

// TRITIUM-PSEUDO fn=residual_add
//| fn residual_add(a, b, n, out):
//|   return n
// TRITIUM-HEIGHT fn=residual_add high=0.00 mid=0.23 low=0.77
int64_t residual_add(double* a, double* b, int64_t n, double* out) {
    int64_t i;

    i = 0L;
    while (i < n) {
        out[i] = (a[i] + b[i]);
        i = (i + 1L);
    }
    return n;
}

// TRITIUM-PSEUDO fn=ftanh
//| fn ftanh(x):
//|   fexp(...)
//|   return x_2 / x_3
// TRITIUM-HEIGHT fn=ftanh high=0.19 mid=0.81 low=0.00 role=leaf_charset
double ftanh(double x) {
    double i;
    int fcmp5;
    int fcmp7;
    double e;
    double x_2;
    double x_3;

    i = x;
    fcmp5 = i > 20.0;
    if (fcmp5) {
        i = 20.0;
    }
    fcmp7 = i < -20.0;
    if (fcmp7) {
        i = -20.0;
    }
    e = fexp(i + i);
    x_2 = e - 1.0;
    x_3 = e + 1.0;
    return x_2 / x_3;
}

// TRITIUM-PSEUDO fn=gelu_tanh
//| fn gelu_tanh(x):
//|   ftanh(...)
//|   return x_5 * x_4
// TRITIUM-HEIGHT fn=gelu_tanh high=0.43 mid=0.57 low=0.00
double gelu_tanh(double x) {
    double x_5;
    double x_2;
    double x_3;
    double x_4;

    x_5 = 0.5 * x;
    x_2 = 0.044715 * ((x * x) * x);
    x_3 = 0.79788456 * (x + x_2);
    x_4 = 1.0 + (ftanh(x_3));
    return x_5 * x_4;
}

// 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.46 mid=0.15 low=0.39
int64_t mlp_gelu_vec(double* x, double* w1, double* w2, int64_t d, int64_t hidden, double* hidden_out, double* out) {
    int64_t i;

    gemv_f(w1, x, hidden, d, hidden_out);
    i = 0L;
    while (i < hidden) {
        hidden_out[i] = (gelu_tanh(hidden_out[i]));
        i = (i + 1L);
    }
    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(...)
//|   gemv_f(...)
//|   residual_add(...)
//|   mlp_gelu_vec(...)
//|   return n_tokens
// TRITIUM-HEIGHT fn=transformer_block high=0.47 mid=0.18 low=0.34
int64_t transformer_block(double* x, int64_t n_tokens, int64_t d, int64_t heads, int64_t dh, int64_t hidden, double scale, double* wq, double* wk, double* wv, double* wo, double* w1, double* w2, double* q, double* k, double* v, double* qh, double* kh, double* vh, double* scores, double* attention, double* flat, double* normed, double* normed2, double* tmp, double* nrm, double* row, double* hidden_out, double* mlp_out, double* projected, double* residual, double* out) {
    int64_t t;
    int64_t i;

    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 = 0L;
    while (t < n_tokens) {
        i = 0L;
        while (i < d) {
            tmp[i] = flat[((t * d) + i)];
            i = (i + 1L);
        }
        gemv_f(wo, tmp, d, d, row);
        i = 0L;
        while (i < d) {
            projected[((t * d) + i)] = row[i];
            i = (i + 1L);
        }
        t = (t + 1L);
    }
    residual_add(x, projected, n_tokens * d, residual);
    ln_rows(residual, n_tokens, d, normed2, tmp, nrm, 0.0);
    t = 0L;
    while (t < n_tokens) {
        i = 0L;
        while (i < d) {
            tmp[i] = normed2[((t * d) + i)];
            i = (i + 1L);
        }
        mlp_gelu_vec(tmp, w1, w2, d, hidden, hidden_out, row);
        i = 0L;
        while (i < d) {
            mlp_out[((t * d) + i)] = row[i];
            i = (i + 1L);
        }
        t = (t + 1L);
    }
    residual_add(residual, mlp_out, n_tokens * d, out);
    return n_tokens;
}