Function with_parsed_tensor_program 

pub fn with_parsed_tensor_program<R>(
    input: &str,
    f: impl for<'a> FnOnce(Result<(&'a [Stmt<'a>], &'a TypeRegistry, PolicyRegistry), String>, &'a Interner, &HashMap<Symbol, (i64, i64)>) -> R,
) -> R

with_parsed_program that also hands over every tensor’s (rows, cols) — the shape table the parser keeps and the desugared statements do not (a declaration desugars to a flat Seq of rows · cols elements). The spatial walk reads it to declare each tensor in the shape the program gives it.