Function test_rope

tests/test_vit.cpp:123–223

Forward-declare the internal function (defined in sam3.cpp) We'll test it by replicating the computation here.

// Forward-declare the internal function (defined in sam3.cpp)
// We'll test it by replicating the computation here.
static bool test_rope(const std::string & ref_dir) {
    fprintf(stderr, "\n=== Test: RoPE ===\n");

    // Test window RoPE (24x24)
    auto ref_window = load_ref(ref_dir + "/rope_window_real");
    if (!ref_window.data.empty()) {
        // ref is [576, 32, 2] — (cos, sin) pairs
        // Our compute_axial_cis produces [N, head_dim] with interleaved (cos, sin)
        // But the Python ref stores them as [N, half_head, 2] where [:,:,0]=cos, [:,:,1]=sin
        const int N = 576;
        const int half_head = 32;
        const int dim = 64;
        const float theta = 10000.0f;
        const float scale_pos = 1.0f;

        std::vector<float> our_rope(N * dim);

        const int half_dim = dim / 4;  // 16
        std::vector<float> freqs(half_dim);
        for (int i = 0; i < half_dim; ++i) {
            freqs[i] = 1.0f / powf(theta, (float)(i * 4) / dim);
        }

        for (int idx = 0; idx < N; ++idx) {
            float t_x = (float)(idx % 24) * scale_pos;
            float t_y = (float)(idx / 24) * scale_pos;

            for (int i = 0; i < half_dim; ++i) {
                float angle_x = t_x * freqs[i];
                our_rope[idx * dim + i * 2 + 0] = cosf(angle_x);
                our_rope[idx * dim + i * 2 + 1] = sinf(angle_x);
            }
            for (int i = 0; i < half_dim; ++i) {
                float angle_y = t_y * freqs[i];
                our_rope[idx * dim + half_dim * 2 + i * 2 + 0] = cosf(angle_y);
                our_rope[idx * dim + half_dim * 2 + i * 2 + 1] = sinf(angle_y);
            }
        }

        // Compare: ref is [N, 32, 2] = [N, half_head, (cos,sin)]
        // Our layout is [N, dim] = [N, 64] with pairs (cos_x, sin_x, cos_x, sin_x, ..., cos_y, sin_y, ...)
        // Need to rearrange comparison:
        // ref[idx, j, 0] = cos of freq j for position idx
        // ref[idx, j, 1] = sin of freq j for position idx
        // Python: freqs_cis = cat([freqs_cis_x, freqs_cis_y], dim=-1) → [N, 32]
        // freqs_cis_x has 16 complex values, freqs_cis_y has 16 complex values → total 32
        // view_as_real → [N, 32, 2] where [:,i,0]=cos, [:,i,1]=sin
        // So ref[idx, j, 0] for j<16 = cos(t_x * freqs[j])
        //    ref[idx, j, 1] for j<16 = sin(t_x * freqs[j])
        //    ref[idx, j, 0] for j>=16 = cos(t_y * freqs[j-16])
        //    ref[idx, j, 1] for j>=16 = sin(t_y * freqs[j-16])

        // Our layout: [idx * 64 + j*2+0] = cos, [idx * 64 + j*2+1] = sin
        // For x: j = 0..15 → our[idx*64 + j*2], ref[idx*64 + j*2]
        // For y: j = 16..31 → our[idx*64 + 32 + (j-16)*2], ref[idx*64 + (j)*2]

        // Actually both are the same layout! ref is stored flat as [N*32*2] = [N*64]
        // and our_rope is [N*64] with the same interleaving.