Steboss

"""
The trunk grid attention passes tokamax a `(N,1,1,N)` boolean key mask built as
the outer product of the token mask (evoformer.py:115). Because the token mask
is a strict prefix (features.py pads with zeros at the end), that boolean mask
carries no more information than one integer per row:

    k_end[b] = L  if token b is real, else 0          (L = number of real tokens)

tokamax accepts exactly that form as `key_value_seq_lengths`, which becomes
`Mask(k_end=...)` and turns the kernel's per-k-tile mask load + `jnp.where` into

Steboss

"""An optimization is valid ONLY if AlphaFold3 still predicts
the same structure and confidence after the change. 

Usage
-----
1. Capture the baseline reference ONCE, from the unmodified inference path:
       python consistency.py capture --af3-output-dir <baseline_out> \
           --ref out/reference.json
2. After each optimization, compare the candidate's output against it:
       python consistency.py check  --af3-output-dir <candidate_out> \

Steboss

"""An optimization is valid ONLY if AlphaFold3 still predicts
the same structure and confidence after the change. 

Usage
-----
1. Capture the baseline reference ONCE, from the unmodified inference path:
       python consistency.py capture --af3-output-dir <baseline_out> \
           --ref out/reference.json
2. After each optimization, compare the candidate's output against it:
       python consistency.py check  --af3-output-dir <candidate_out> \

Steboss

"""An optimization is valid ONLY if AlphaFold3 still predicts
the same structure and confidence after the change. 

Usage
-----
1. Capture the baseline reference ONCE, from the unmodified inference path:
       python consistency.py capture --af3-output-dir <baseline_out> \
           --ref out/reference.json
2. After each optimization, compare the candidate's output against it:
       python consistency.py check  --af3-output-dir <candidate_out> \

Steboss

Steboss

Steboss

Claim, Why isn't true, Correct claim
"It fits because the file fits", "The file is the weight payload at one precision. Peak memory also includes activations, workspace, and a KV cache that grows linearly with the context", "Fit is undetermined wihtout context length and runtime overhead"

Steboss

Claim, Why isn't true, Correct claim
"It fits because the file fits", "The file is the weight payload at one precision. Peak memory also includes activations, workspace, and a KV cache that grows linearly with the context", "Fit is undetermined wihtout context length and runtime overhead"

Steboss

use rayon::prelude::*;
use std::arch::aarch64::*;
use std::env;
use std::fs;

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() != 4 {
        eprintln!("Usage: {} M K N", args[0]);
        return;

Steboss

use rayon::prelude::*;
use std::env;

const N: usize = 2; // Fixed block size

// HELPER: Convert slice to fixed array for safety (optional but good for debugging)
fn to_array(slice: &[f64]) -> [f64; N * N] {
    slice.try_into().expect("Slice has incorrect length")
}