/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * Unit tests for Foundry Ternarizer module.
 */

#include "../../lib/test.h"
#include "../ternarize.h"
#include "../ft_layer_ops.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <sys/stat.h>

/* ============================================================= */
/* Config tests                                                  */
/* ============================================================= */

static void test_config_defaults(void) {
    ft_config cfg = ft_config_defaults();
    ASSERT_INT_EQ(128, cfg.group_size);
    ASSERT_FLT_NEAR(0.00008f, cfg.k_frac_start, 1e-7f);
    ASSERT_FLT_NEAR(0.000005f, cfg.k_frac_end, 1e-8f);
    ASSERT_FLT_NEAR(0.3f, cfg.snr_weight, 1e-6f);
    ASSERT_FLT_NEAR(0.15f, cfg.hess_power, 1e-6f);
    ASSERT_FLT_NEAR(0.2f, cfg.momentum_bonus, 1e-6f);
    ASSERT_FLT_NEAR(0.17f, cfg.boltz_frac_start, 1e-6f);
    ASSERT_FLT_NEAR(0.0f, cfg.boltz_frac_end, 1e-6f);
    ASSERT_FLT_NEAR(1.8f, cfg.temporal_weight, 1e-6f);
    ASSERT_INT_EQ(60, cfg.phase1_iters);
    ASSERT_INT_EQ(140, cfg.phase2_iters);
    ASSERT_INT_EQ(128, cfg.batch_size);
}

static void test_arch_names(void) {
    ASSERT_STR_EQ("Llama", ft_arch_name(FT_ARCH_LLAMA));
    ASSERT_STR_EQ("Qwen2", ft_arch_name(FT_ARCH_QWEN2));
    ASSERT_STR_EQ("Qwen3", ft_arch_name(FT_ARCH_QWEN3));
    ASSERT_STR_EQ("Phi-3", ft_arch_name(FT_ARCH_PHI3));
    ASSERT_STR_EQ("Gemma2", ft_arch_name(FT_ARCH_GEMMA2));
    ASSERT_STR_EQ("RWKV-X", ft_arch_name(FT_ARCH_RWKVX));
    ASSERT_STR_EQ("Unknown", ft_arch_name(FT_ARCH_UNKNOWN));
}

/* ============================================================= */
/* Ternary bit manipulation tests                                */
/* ============================================================= */

static void test_ternary_get_set(void) {
    /* Create a minimal layer state */
    ft_config cfg = ft_config_defaults();
    cfg.n_layers = 1;
    cfg.n_embd = 64;
    cfg.n_heads = 4;
    cfg.n_kv_heads = 4;
    cfg.head_dim = 16;
    cfg.ffn_dim = 128;

    /* Allocate manually for test */
    int64_t numel = 64; /* Small test */
    ft_layer_state s;
    memset(&s, 0, sizeof(s));
    s.numel_total = numel;
    s.group_size = 32;
    int64_t n_bytes = (numel + 7) / 8;
    s.sign = calloc(n_bytes, 1);
    s.mask = calloc(n_bytes, 1);
    s.momentum = calloc(n_bytes, 1);

    /* Test set/get for all ternary values */
    ft_set_ternary(&s, 1, 0);
    ASSERT_INT_EQ(1, ft_get_ternary(&s, 0));

    ft_set_ternary(&s, -1, 1);
    ASSERT_INT_EQ(-1, ft_get_ternary(&s, 1));

    ft_set_ternary(&s, 0, 2);
    ASSERT_INT_EQ(0, ft_get_ternary(&s, 2));

    /* Test across byte boundary */
    ft_set_ternary(&s, 1, 7);
    ft_set_ternary(&s, -1, 8);
    ft_set_ternary(&s, 1, 15);
    ft_set_ternary(&s, -1, 16);

    ASSERT_INT_EQ(1, ft_get_ternary(&s, 7));
    ASSERT_INT_EQ(-1, ft_get_ternary(&s, 8));
    ASSERT_INT_EQ(1, ft_get_ternary(&s, 15));
    ASSERT_INT_EQ(-1, ft_get_ternary(&s, 16));

    /* Overwrite and verify */
    ft_set_ternary(&s, 0, 0);
    ASSERT_INT_EQ(0, ft_get_ternary(&s, 0));
    ft_set_ternary(&s, -1, 0);
    ASSERT_INT_EQ(-1, ft_get_ternary(&s, 0));

    /* Verify adjacent values weren't corrupted */
    ASSERT_INT_EQ(-1, ft_get_ternary(&s, 1));
    ASSERT_INT_EQ(0, ft_get_ternary(&s, 2));

    free(s.sign);
    free(s.mask);
    free(s.momentum);
}

static void test_momentum_get_set(void) {
    ft_layer_state s;
    memset(&s, 0, sizeof(s));
    s.numel_total = 32;
    int64_t n_bytes = 4;
    s.momentum = calloc(n_bytes, 1);

    ft_set_momentum(&s, 1, 0);
    ASSERT_INT_EQ(1, ft_get_momentum(&s, 0));

    ft_set_momentum(&s, 0, 0);
    ASSERT_INT_EQ(0, ft_get_momentum(&s, 0));

    ft_set_momentum(&s, 1, 7);
    ft_set_momentum(&s, 1, 8);
    ASSERT_INT_EQ(1, ft_get_momentum(&s, 7));
    ASSERT_INT_EQ(1, ft_get_momentum(&s, 8));
    ASSERT_INT_EQ(0, ft_get_momentum(&s, 6));

    free(s.momentum);
}

/* ============================================================= */
/* RMSNorm tests                                                 */
/* ============================================================= */

static void test_rmsnorm_basic(void) {
    float x[] = {1.0f, 2.0f, 3.0f, 4.0f};
    float w[] = {1.0f, 1.0f, 1.0f, 1.0f};
    float out[4] = {0};

    ft_rmsnorm(x, w, out, 4, 1e-6f);

    /* rms = sqrt((1+4+9+16)/4) = sqrt(7.5) ≈ 2.7386 */
    float rms = sqrtf(7.5f + 1e-6f);
    ASSERT_FLT_NEAR(1.0f / rms, out[0], 0.001f);
    ASSERT_FLT_NEAR(2.0f / rms, out[1], 0.001f);
    ASSERT_FLT_NEAR(3.0f / rms, out[2], 0.001f);
    ASSERT_FLT_NEAR(4.0f / rms, out[3], 0.001f);
}

static void test_rmsnorm_with_weight(void) {
    float x[] = {1.0f, -1.0f};
    float w[] = {2.0f, 3.0f};
    float out[2] = {0};

    ft_rmsnorm(x, w, out, 2, 1e-6f);

    float rms = sqrtf(1.0f + 1e-6f);
    ASSERT_FLT_NEAR(2.0f / rms, out[0], 0.001f);
    ASSERT_FLT_NEAR(-3.0f / rms, out[1], 0.001f);
}

/* ============================================================= */
/* SiLU tests                                                    */
/* ============================================================= */

static void test_silu_basic(void) {
    float x[] = {0.0f, 1.0f, -1.0f, 5.0f};
    float out[4] = {0};

    ft_silu(x, out, 4);

    /* silu(0) = 0 * 0.5 = 0 */
    ASSERT_FLT_NEAR(0.0f, out[0], 0.001f);

    /* silu(1) = 1 * sigmoid(1) ≈ 1 * 0.7311 = 0.7311 */
    float s1 = 1.0f / (1.0f + expf(-1.0f));
    ASSERT_FLT_NEAR(1.0f * s1, out[1], 0.001f);

    /* silu(-1) = -1 * sigmoid(-1) ≈ -1 * 0.2689 = -0.2689 */
    float sm1 = 1.0f / (1.0f + expf(1.0f));
    ASSERT_FLT_NEAR(-1.0f * sm1, out[2], 0.001f);
}

/* ============================================================= */
/* Ternary matvec tests                                          */
/* ============================================================= */

static void test_ternary_matvec_identity(void) {
    /* 2x2 identity matrix in ternary: [[1,0],[0,1]] */
    ft_layer_state s;
    memset(&s, 0, sizeof(s));
    s.numel_total = 4;
    s.group_size = 4;
    s.n_groups = 1;
    int64_t n_bytes = 1;
    s.sign = calloc(n_bytes, 1);
    s.mask = calloc(n_bytes, 1);
    s.momentum = calloc(n_bytes, 1);
    s.scales = calloc(1, sizeof(float));
    s.scales[0] = 1.0f;

    /* Set identity: [1,0; 0,1] */
    /* (0,0) = +1, (0,1) = 0, (1,0) = 0, (1,1) = +1 */
    ft_set_ternary(&s, 1, 0);  /* [0][0] */
    ft_set_ternary(&s, 0, 1);  /* [0][1] */
    ft_set_ternary(&s, 0, 2);  /* [1][0] */
    ft_set_ternary(&s, 1, 3);  /* [1][1] */

    float x[] = {3.0f, 7.0f};
    float out[2] = {0};
    ft_ternary_matvec(&s, 0, 2, 2, x, out);

    ASSERT_FLT_NEAR(3.0f, out[0], 0.001f);
    ASSERT_FLT_NEAR(7.0f, out[1], 0.001f);

    free(s.sign); free(s.mask); free(s.momentum); free(s.scales);
}

static void test_ternary_matvec_negation(void) {
    /* 2x2 negation matrix: [[-1,0],[0,-1]] */
    ft_layer_state s;
    memset(&s, 0, sizeof(s));
    s.numel_total = 4;
    s.group_size = 4;
    s.n_groups = 1;
    int64_t n_bytes = 1;
    s.sign = calloc(n_bytes, 1);
    s.mask = calloc(n_bytes, 1);
    s.momentum = calloc(n_bytes, 1);
    s.scales = calloc(1, sizeof(float));
    s.scales[0] = 1.0f;

    ft_set_ternary(&s, -1, 0);
    ft_set_ternary(&s, 0, 1);
    ft_set_ternary(&s, 0, 2);
    ft_set_ternary(&s, -1, 3);

    float x[] = {3.0f, 7.0f};
    float out[2] = {0};
    ft_ternary_matvec(&s, 0, 2, 2, x, out);

    ASSERT_FLT_NEAR(-3.0f, out[0], 0.001f);
    ASSERT_FLT_NEAR(-7.0f, out[1], 0.001f);

    free(s.sign); free(s.mask); free(s.momentum); free(s.scales);
}

static void test_ternary_matvec_with_scale(void) {
    /* 2x2 identity with scale=2.5 */
    ft_layer_state s;
    memset(&s, 0, sizeof(s));
    s.numel_total = 4;
    s.group_size = 4;
    s.n_groups = 1;
    int64_t n_bytes = 1;
    s.sign = calloc(n_bytes, 1);
    s.mask = calloc(n_bytes, 1);
    s.momentum = calloc(n_bytes, 1);
    s.scales = calloc(1, sizeof(float));
    s.scales[0] = 2.5f;

    ft_set_ternary(&s, 1, 0);
    ft_set_ternary(&s, 0, 1);
    ft_set_ternary(&s, 0, 2);
    ft_set_ternary(&s, 1, 3);

    float x[] = {1.0f, 2.0f};
    float out[2] = {0};
    ft_ternary_matvec(&s, 0, 2, 2, x, out);

    ASSERT_FLT_NEAR(2.5f, out[0], 0.001f);
    ASSERT_FLT_NEAR(5.0f, out[1], 0.001f);

    free(s.sign); free(s.mask); free(s.momentum); free(s.scales);
}

/* ============================================================= */
/* Ternary matvec backward tests                                 */
/* ============================================================= */

static void test_ternary_matvec_backward_input(void) {
    /* Identity backward: d_x should equal d_out */
    ft_layer_state s;
    memset(&s, 0, sizeof(s));
    s.numel_total = 4;
    s.group_size = 4;
    s.n_groups = 1;
    int64_t n_bytes = 1;
    s.sign = calloc(n_bytes, 1);
    s.mask = calloc(n_bytes, 1);
    s.momentum = calloc(n_bytes, 1);
    s.scales = calloc(1, sizeof(float));
    s.scales[0] = 1.0f;

    ft_set_ternary(&s, 1, 0);
    ft_set_ternary(&s, 0, 1);
    ft_set_ternary(&s, 0, 2);
    ft_set_ternary(&s, 1, 3);

    float d_out[] = {5.0f, 3.0f};
    float d_x[2] = {0};
    ft_ternary_matvec_backward_input(&s, 0, 2, 2, d_out, d_x);

    ASSERT_FLT_NEAR(5.0f, d_x[0], 0.001f);
    ASSERT_FLT_NEAR(3.0f, d_x[1], 0.001f);

    free(s.sign); free(s.mask); free(s.momentum); free(s.scales);
}

static void test_ternary_matvec_backward_weight(void) {
    /* Verify gradient accumulation */
    ft_layer_state s;
    memset(&s, 0, sizeof(s));
    s.numel_total = 4;
    s.group_size = 4;
    s.n_groups = 1;
    int64_t n_bytes = 1;
    s.sign = calloc(n_bytes, 1);
    s.mask = calloc(n_bytes, 1);
    s.momentum = calloc(n_bytes, 1);
    s.scales = calloc(1, sizeof(float));
    s.scales[0] = 1.0f;

    ft_set_ternary(&s, 1, 0);
    ft_set_ternary(&s, 0, 1);
    ft_set_ternary(&s, 0, 2);
    ft_set_ternary(&s, 1, 3);

    ft_grad_stats stats;
    memset(&stats, 0, sizeof(stats));
    stats.numel = 4;
    stats.grad_sum = calloc(4, sizeof(float));
    stats.grad_sq_sum = calloc(4, sizeof(float));
    stats.hessian_diag = calloc(4, sizeof(float));

    float x[] = {2.0f, 3.0f};
    float d_out[] = {5.0f, 7.0f};

    ft_ternary_matvec_backward_weight(&s, 0, 2, 2, x, d_out, &stats);

    /* grad[0,0] = x[0]*d_out[0] = 2*5 = 10 */
    /* grad[0,1] = x[1]*d_out[0] = 3*5 = 15 */
    /* grad[1,0] = x[0]*d_out[1] = 2*7 = 14 */
    /* grad[1,1] = x[1]*d_out[1] = 3*7 = 21 */
    ASSERT_FLT_NEAR(10.0f, stats.grad_sum[0], 0.001f);
    ASSERT_FLT_NEAR(15.0f, stats.grad_sum[1], 0.001f);
    ASSERT_FLT_NEAR(14.0f, stats.grad_sum[2], 0.001f);
    ASSERT_FLT_NEAR(21.0f, stats.grad_sum[3], 0.001f);

    /* grad_sq: 100, 225, 196, 441 */
    ASSERT_FLT_NEAR(100.0f, stats.grad_sq_sum[0], 0.01f);
    ASSERT_FLT_NEAR(225.0f, stats.grad_sq_sum[1], 0.01f);

    free(stats.grad_sum);
    free(stats.grad_sq_sum);
    free(stats.hessian_diag);
    free(s.sign); free(s.mask); free(s.momentum); free(s.scales);
}

/* ============================================================= */
/* Layer state I/O tests                                         */
/* ============================================================= */

static void test_layer_save_load(void) {
    ft_config cfg = ft_config_defaults();
    cfg.n_layers = 1;
    cfg.n_embd = 32;
    cfg.n_heads = 2;
    cfg.n_kv_heads = 2;
    cfg.head_dim = 16;
    cfg.ffn_dim = 64;
    snprintf(cfg.workdir, sizeof(cfg.workdir), "/tmp/ft_test_io");

    /* Create and populate state */
    ft_layer_state *s = calloc(1, sizeof(ft_layer_state));
    s->layer_idx = 0;
    s->numel_total = 32 * 32 + 32 * 32 + 32 * 32 + 32 * 32 + 64 * 32 + 64 * 32 + 32 * 64;
    /* Simplified: just use a small number for test */
    s->numel_total = 64;
    s->group_size = cfg.group_size;
    s->n_groups = (int)((s->numel_total + cfg.group_size - 1) / cfg.group_size);

    int64_t n_bytes = (s->numel_total + 7) / 8;
    s->sign = calloc(n_bytes, 1);
    s->mask = calloc(n_bytes, 1);
    s->momentum = calloc(n_bytes, 1);
    s->scales = calloc(s->n_groups, sizeof(float));
    s->prev_scores = calloc(s->numel_total, sizeof(float));
    s->norm_attn = calloc(cfg.n_embd, sizeof(float));
    s->norm_ffn = calloc(cfg.n_embd, sizeof(float));
    s->weights = NULL;
    s->n_weights = 0;

    /* Set some values */
    ft_set_ternary(s, 1, 0);
    ft_set_ternary(s, -1, 1);
    ft_set_ternary(s, 0, 2);
    ft_set_ternary(s, 1, 63);
    s->scales[0] = 1.5f;
    s->iteration = 42;
    s->phase = 1;
    s->total_flips = 1234;
    s->best_loss = 0.567f;
    s->norm_attn[0] = 1.0f;
    s->norm_ffn[0] = 2.0f;

    /* Save */
    ASSERT_TRUE(ft_layer_save(s, &cfg));

    /* Reload into a fresh state */
    /* ft_layer_load expects proper config for numel computation,
     * but our test numel doesn't match. So we test the I/O directly. */
    char path[1024];
    snprintf(path, sizeof(path), "%s/layer_000/ternary.bin", cfg.workdir);
    FILE *f = fopen(path, "rb");
    ASSERT_NOT_NULL(f);

    int64_t stored_numel;
    int iter, phase;
    int64_t total_flips;
    float best_loss;
    fread(&stored_numel, sizeof(int64_t), 1, f);
    fread(&iter, sizeof(int), 1, f);
    fread(&phase, sizeof(int), 1, f);
    fread(&total_flips, sizeof(int64_t), 1, f);
    fread(&best_loss, sizeof(float), 1, f);

    ASSERT_INT_EQ(64, (int)stored_numel);
    ASSERT_INT_EQ(42, iter);
    ASSERT_INT_EQ(1, phase);
    ASSERT_INT_EQ(1234, (int)total_flips);
    ASSERT_FLT_NEAR(0.567f, best_loss, 0.001f);

    /* Read back sign/mask/momentum and verify ternary values */
    uint8_t *sign2 = calloc(n_bytes, 1);
    uint8_t *mask2 = calloc(n_bytes, 1);
    uint8_t *mom2 = calloc(n_bytes, 1);
    fread(sign2, 1, n_bytes, f);
    fread(mask2, 1, n_bytes, f);
    fread(mom2, 1, n_bytes, f);
    fclose(f);

    /* Verify bit patterns match */
    ASSERT_TRUE(memcmp(s->sign, sign2, n_bytes) == 0);
    ASSERT_TRUE(memcmp(s->mask, mask2, n_bytes) == 0);

    free(sign2); free(mask2); free(mom2);
    free(s->sign); free(s->mask); free(s->momentum);
    free(s->scales); free(s->prev_scores);
    free(s->norm_attn); free(s->norm_ffn);
    free(s);

    /* Cleanup */
    system("rm -rf /tmp/ft_test_io");
}

/* ============================================================= */
/* Grad stats I/O tests                                          */
/* ============================================================= */

static void test_grad_stats_save_load(void) {
    ft_config cfg = ft_config_defaults();
    snprintf(cfg.workdir, sizeof(cfg.workdir), "/tmp/ft_test_grad");

    /* Create workdir */
    system("mkdir -p /tmp/ft_test_grad/layer_000");

    ft_grad_stats g;
    g.numel = 16;
    g.n_samples = 42;
    g.grad_sum = malloc(16 * sizeof(float));
    g.grad_sq_sum = malloc(16 * sizeof(float));
    g.hessian_diag = malloc(16 * sizeof(float));

    for (int i = 0; i < 16; i++) {
        g.grad_sum[i] = (float)i * 0.1f;
        g.grad_sq_sum[i] = (float)i * 0.01f;
        g.hessian_diag[i] = (float)i * 0.001f;
    }

    ASSERT_TRUE(ft_grad_stats_save(&g, &cfg, 0));

    ft_grad_stats *g2 = ft_grad_stats_load(&cfg, 0);
    ASSERT_NOT_NULL(g2);
    ASSERT_INT_EQ(16, (int)g2->numel);
    ASSERT_INT_EQ(42, g2->n_samples);

    for (int i = 0; i < 16; i++) {
        ASSERT_FLT_NEAR(g.grad_sum[i], g2->grad_sum[i], 1e-6f);
        ASSERT_FLT_NEAR(g.grad_sq_sum[i], g2->grad_sq_sum[i], 1e-6f);
        ASSERT_FLT_NEAR(g.hessian_diag[i], g2->hessian_diag[i], 1e-6f);
    }

    free(g.grad_sum); free(g.grad_sq_sum); free(g.hessian_diag);
    ft_grad_stats_free(g2);
    system("rm -rf /tmp/ft_test_grad");
}

/* ============================================================= */
/* Status I/O tests                                              */
/* ============================================================= */

static void test_status_io(void) {
    ft_config cfg = ft_config_defaults();
    snprintf(cfg.workdir, sizeof(cfg.workdir), "/tmp/ft_test_status");
    system("mkdir -p /tmp/ft_test_status");

    ft_status st;
    memset(&st, 0, sizeof(st));
    st.n_layers = 32;
    st.layers_split = 32;
    st.layers_swept = 16;
    st.layers_flipped = 8;
    st.current_sweep_iter = 3;
    st.overall_loss = 2.345f;
    snprintf(st.model_name, sizeof(st.model_name), "test-model");

    ASSERT_TRUE(ft_status_write(&st, &cfg));

    ft_status st2 = ft_status_read(&cfg);
    ASSERT_INT_EQ(32, st2.n_layers);
    ASSERT_INT_EQ(32, st2.layers_split);
    ASSERT_INT_EQ(16, st2.layers_swept);
    ASSERT_INT_EQ(8, st2.layers_flipped);
    ASSERT_INT_EQ(3, st2.current_sweep_iter);
    ASSERT_FLT_NEAR(2.345f, st2.overall_loss, 0.01f);
    ASSERT_STR_EQ("test-model", st2.model_name);

    system("rm -rf /tmp/ft_test_status");
}

/* ============================================================= */
/* SiLU backward test                                            */
/* ============================================================= */

static void test_silu_backward(void) {
    float x[] = {0.0f, 1.0f, -1.0f};
    float d_out[] = {1.0f, 1.0f, 1.0f};
    float d_x[3] = {0};

    ft_silu_backward(x, d_out, d_x, 3);

    /* d_silu(0)/dx = sigmoid(0) + 0 * sigmoid(0) * (1-sigmoid(0)) = 0.5 */
    ASSERT_FLT_NEAR(0.5f, d_x[0], 0.01f);

    /* d_silu(1)/dx = sigmoid(1) + 1 * sigmoid(1) * (1-sigmoid(1)) */
    float s1 = 1.0f / (1.0f + expf(-1.0f));
    float expected = s1 + 1.0f * s1 * (1.0f - s1);
    ASSERT_FLT_NEAR(expected, d_x[1], 0.01f);
}

/* ============================================================= */
/* RMSNorm backward test                                         */
/* ============================================================= */

static void test_rmsnorm_backward(void) {
    /* Numerical gradient check */
    float x[] = {1.0f, 2.0f, 3.0f, 4.0f};
    float w[] = {1.0f, 1.0f, 1.0f, 1.0f};
    float d_out[] = {1.0f, 0.0f, 0.0f, 0.0f};
    float d_x[4] = {0};

    ft_rmsnorm_backward(x, w, d_out, d_x, 4, 1e-6f);

    /* Numerical gradient for d_x[0] */
    float eps = 1e-3f;
    float x_plus[4] = {1.0f + eps, 2.0f, 3.0f, 4.0f};
    float x_minus[4] = {1.0f - eps, 2.0f, 3.0f, 4.0f};
    float out_plus[4], out_minus[4];
    ft_rmsnorm(x_plus, w, out_plus, 4, 1e-6f);
    ft_rmsnorm(x_minus, w, out_minus, 4, 1e-6f);

    /* d_out = [1,0,0,0], so numeric grad = (out_plus[0] - out_minus[0]) / (2*eps) */
    float numeric = (out_plus[0] - out_minus[0]) / (2.0f * eps);
    ASSERT_FLT_NEAR(numeric, d_x[0], 0.01f);
}

/* ============================================================= */
/* Main                                                          */
/* ============================================================= */

int main(void) {
    printf("\n  Running ternarize tests...\n\n");

    /* Config */
    RUN_TEST(config_defaults);
    RUN_TEST(arch_names);

    /* Ternary bit manipulation */
    RUN_TEST(ternary_get_set);
    RUN_TEST(momentum_get_set);

    /* Layer ops */
    RUN_TEST(rmsnorm_basic);
    RUN_TEST(rmsnorm_with_weight);
    RUN_TEST(rmsnorm_backward);
    RUN_TEST(silu_basic);
    RUN_TEST(silu_backward);

    /* Ternary matvec */
    RUN_TEST(ternary_matvec_identity);
    RUN_TEST(ternary_matvec_negation);
    RUN_TEST(ternary_matvec_with_scale);
    RUN_TEST(ternary_matvec_backward_input);
    RUN_TEST(ternary_matvec_backward_weight);

    /* I/O */
    RUN_TEST(layer_save_load);
    RUN_TEST(grad_stats_save_load);
    RUN_TEST(status_io);

    TEST_REPORT();
    return __ft_failed > 0 ? 1 : 0;
}