/**
 * model_cpu.ts — the exact CPU forward + backward for a whole
 * {@link HybridMambaModel}: token embedding → layer stack → final RMSNorm →
 * tied LM head → cross-entropy.
 *
 * This is the gradient engine {@link MambaTrainer} uses. Before it existed the
 * trainer handed the SAME `dLogits` buffer to AdamW for every parameter, so no
 * parameter ever received its own gradient — `learn()` was safe but inert. Now
 * the backward pass is real, and `tests/backprop.test.ts` proves it against
 * finite differences without needing a GPU.
 *
 * Parameter names match {@link HybridMambaModel.parameters} exactly
 * (`embedding`, `layer{i}.{name}`, `final_norm`), which is how the trainer
 * pairs each parameter with its OWN gradient.
 */
import { type Mamba1CpuWeights, type Mamba1CpuGrads } from './mamba1_cpu.js';
/** Layer types the CPU gradient engine can differentiate today. */
export declare const CPU_GRADIENT_LAYER_TYPES: readonly ["mamba1"];
export type CpuGradientLayerType = (typeof CPU_GRADIENT_LAYER_TYPES)[number];
export interface CpuModelDims {
    vocabSize: number;
    dModel: number;
    dState: number;
    dConv: number;
    dInner: number;
    dtRank: number;
    numLayers: number;
}
/** Every tensor the CPU model reads, mirrored out of the GPU buffers. */
export interface CpuModelWeights {
    embedding: Float32Array;
    finalNorm: Float32Array;
    lmHeadBias: Float32Array;
    layers: Mamba1CpuWeights[];
}
export interface CpuModelGrads {
    embedding: Float32Array;
    finalNorm: Float32Array;
    lmHeadBias: Float32Array;
    layers: Mamba1CpuGrads[];
}
/** Knobs on one backward pass. */
export interface CpuBackwardOptions {
    /**
     * Recompute each layer's activations during the backward pass instead of
     * retaining every layer's cache from the forward. Trades one extra forward per
     * layer for O(1) instead of O(layers) retained activations; the gradients are
     * numerically identical (`tests/peft_trainer.test.ts` asserts bit equality).
     */
    activationCheckpointing?: boolean;
}
/** Loss + the gradient of every parameter, keyed by `parameters()` name. */
export interface CpuBackwardResult {
    loss: number;
    grads: Map<string, Float32Array>;
}
export declare function zeroCpuModelGrads(dims: CpuModelDims): CpuModelGrads;
/** Flatten a {@link CpuModelGrads} into the `parameters()` naming the optimizer uses. */
export declare function toNamedGrads(g: CpuModelGrads): Map<string, Float32Array>;
/** The model's tensors keyed the same way — the inverse of {@link cpuWeightsFromNamed}. */
export declare function namedWeights(w: CpuModelWeights): Map<string, Float32Array>;
/**
 * Assemble a {@link CpuModelWeights} from tensors keyed by `parameters()` name —
 * how {@link MambaTrainer} mirrors the GPU buffers into the gradient engine.
 * Throws on a missing or wrongly-sized tensor rather than training on garbage.
 */
export declare function cpuWeightsFromNamed(named: Map<string, Float32Array>, dims: CpuModelDims): CpuModelWeights;
/** Logits only — the CPU mirror of `HybridMambaModel.forward`. */
export declare function cpuModelForward(tokenIds: ArrayLike<number>, w: CpuModelWeights, dims: CpuModelDims, batch: number, seqLen: number): Float32Array;
/** Mean next-token cross-entropy over the sequence — the trainer's loss. */
export declare function cpuModelLoss(tokenIds: ArrayLike<number>, targets: ArrayLike<number>, w: CpuModelWeights, dims: CpuModelDims, batch: number, seqLen: number): number;
/**
 * Forward + backward. Returns the mean cross-entropy loss and the gradient of
 * EVERY parameter, keyed by its `parameters()` name.
 *
 * The loss is averaged over the `batch * seqLen` predicted positions, so the
 * gradients are the gradients of that mean — the quantity the trainer reports.
 */
export declare function cpuModelBackward(tokenIds: ArrayLike<number>, targets: ArrayLike<number>, w: CpuModelWeights, dims: CpuModelDims, batch: number, seqLen: number, opts?: CpuBackwardOptions): CpuBackwardResult;
/**
 * The 2-D shape of a parameter, or `null` when it is a vector (bias / norm gain
 * / `D_vec`).
 *
 * Low-rank adaptation only means anything over a matrix, so this is the single
 * place that says which parameters an adapter can cover — and it derives the
 * shapes from `dims` rather than restating the block's layout, so it cannot
 * drift from {@link cpuWeightsFromNamed}.
 */
export declare function paramMatrixShape(name: string, dims: CpuModelDims): [rows: number, cols: number] | null;
/**
 * Parameters a low-rank adapter covers by default: the four projections inside
 * each block.
 *
 * Deliberately NOT the token embedding (by far the largest tensor, and tied to
 * the output head — adapting it low-rank changes the head too) and NOT `A_log`
 * (the state-transition decay, which WSLA freezes precisely because drifting it
 * destabilises the SSM).
 */
export declare const DEFAULT_LORA_TARGETS: readonly ["wInProj", "wXProj", "wDtProj", "wOutProj"];
//# sourceMappingURL=model_cpu.d.ts.map