/**
 * adamw.ts — AdamW optimiser over a model's flat parameter list.
 *
 * Shared by every CPU-reference trainer in the engine (MoE FFN, the full
 * EvermindLM) so the optimiser maths lives in exactly one place. Operates on any
 * object exposing index-aligned `parameters()` / `gradients()` Float32Arrays.
 */
export interface OptimParam {
    data: Float32Array;
}
export interface OptimTarget {
    parameters(): OptimParam[];
    gradients(): OptimParam[];
}
/**
 * Optimizer-state sharding (ZeRO-1 analog). When set, this optimiser instance
 * OWNS only parameter tensors where `paramIndex % count === index`, and only
 * allocates the AdamW moments (m, v — 2× the parameter bytes, the dominant cost
 * of full fine-tuning per the cookbook's memory accounting) for its shard.
 *
 * Run `count` instances over the SAME model, one per `index`, and the union of
 * their steps equals one unsharded step — but each holds only 1/count of the
 * optimizer state. That is exactly how FSDP/ZeRO fits a larger model on the same
 * hardware; here it is the seam a multi-device trainer partitions on, and even
 * single-process it caps resident moment memory.
 */
export interface ShardSpec {
    /** This shard's rank, 0-based. */
    index: number;
    /** Total number of shards. */
    count: number;
}
export interface AdamWOptions {
    lr?: number;
    beta1?: number;
    beta2?: number;
    eps?: number;
    weightDecay?: number;
    /** Optimizer-state sharding. Omit for the full (single-owner) optimiser. */
    shard?: ShardSpec;
}
export declare class AdamW {
    private readonly target;
    private readonly m;
    private readonly v;
    private t;
    private readonly opt;
    private readonly shard;
    constructor(target: OptimTarget, options?: AdamWOptions);
    /** Whether this shard owns (and updates) parameter tensor `p`. */
    private _owns;
    /** One optimiser step from the currently-accumulated gradients (owned tensors only). */
    step(): void;
    /** Bytes of optimizer state this shard holds (m + v) — 1/count of the full state. */
    stateBytes(): number;
}
/** Hyperparameters of one {@link adamwUpdateInPlace} step. */
export interface AdamWKernelStep {
    lr: number;
    beta1: number;
    beta2: number;
    eps: number;
    weightDecay: number;
    /** beta1^t — the precomputed bias-correction term (matches the kernel uniform). */
    beta1_t: number;
    /** beta2^t. */
    beta2_t: number;
    /** Trust region: max |Δθ| per step. 0 disables the bound. */
    maxDelta: number;
    /** Global gradient-clip factor, applied to `grad` on the fly (1 = no clipping). */
    gradScale?: number;
}
/**
 * One AdamW step over a single parameter tensor — the exact CPU mirror of the
 * `adamw_update` entry point in `WEIGHT_UPDATE_WGSL`, including its two safety
 * properties:
 *
 *   • the **NaN/Inf guard** — a non-finite step is dropped rather than written
 *     into a weight, so one bad gradient cannot permanently poison the model;
 *   • the **trust region** — `maxDelta` bounds how far a single step can move a
 *     weight, which is what keeps repeated write-through adapts reversible.
 *
 * Decoupled weight decay (`θ * (1 - lr·wd)`), matching the kernel — deliberately
 * NOT the coupled form {@link AdamW} uses for the MoE/EvermindLM trainers.
 *
 * `param`, `m` and `v` are updated in place.
 */
export declare function adamwUpdateInPlace(param: Float32Array, grad: Float32Array, m: Float32Array, v: Float32Array, hp: AdamWKernelStep): void;
//# sourceMappingURL=adamw.d.ts.map