Module infer 

TabFM inference engine, with a real batch dimension B (predict_batch) — e.g. one batch item per ensemble member, so ensemble::orchestrate can share the fixed cost of the 24-block ICL stage across all members in one forward pass instead of paying it once per member. The single-table predict() is a thin B=1 wrapper around the same code path; none of the attention/RMSNorm/RoPE math below changed to add batching — only the four “stage” functions (cell_embed, col_embedding_forward, row_interaction_forward, icl_forward) gained a leading B dimension, via reshapes around the same 3D attention calls they always made (masks/weights are shared scalars across the batch — every member has the same row/feature count, train_size, and d; only cell values and cat_mask vary per member).

Architecture (from tabfm/src/pytorch/model.py, verified against the installed package): CellEmbedder (per-cell grouped Fourier features + train-row y-embedding) -> ColEmbedding (SetTransformer / induced attention over rows, masked to train rows) -> prepend row_num_cls learned CLS tokens on the column axis -> RowInteraction (RoPE cross-column self-attention, masked to valid/unpadded columns) -> ColEmbedding (stage 2) -> RowInteraction (stage 2, output collapsed to the CLS-token slice -> icl_dim) -> ICLearning (24-block self-attention over rows, y re-injected at train rows, masked so only train rows are attendable keys) -> MLP decoder -> per-class logits or a scalar.

Numeric details that matter for parity (see model-to-gguf skill’s debugging ladder):

• RoPE frequencies are checkpoint-loaded buffers, never recomputed from a formula.
• MultiheadAttention pre-scales q by a learned per-dimension softplus’d scale, then calls attention with scale=1.0 — do not apply an additional 1/sqrt(d).
• All masking is additive key-side masking (no causal masking anywhere).
• RoPE here is the interleaved-pair variant (x[0::2], x[1::2]), NOT the Llama rotate-half variant used elsewhere in this repo (see toto’s infer/rope.rs) — deliberately not reused.

Structs

InferConfig

TabFMModel

TabFMModelBuilder

Fluent constructor for TabFMModel: point it at a GGUF file and a config (from a parsed config.json via config_from, or a hand-built InferConfig via config), then call build.