Interface TrainableGraph

What the trainer consumes from a graph. A graph declares only the parts it implements; the trainer names the missing declaration when it needs one.

interface TrainableGraph {
    denoising?: DenoisingTraining;
    flashAttention?: FlashTrainingAttention;
    gradCheckpoint?: GradCheckpointTraining;
    prefixShared?: PrefixSharedTraining;
    segmented?: SegmentedTraining;
    lmHead?(): HeadQuant;
}

Index

Properties

denoising? flashAttention? gradCheckpoint? prefixShared? segmented?

Methods

Properties

`Optional` `Readonly`denoising

denoising?: DenoisingTraining

`Optional` `Readonly`flashAttention

flashAttention?: FlashTrainingAttention

`Optional` `Readonly`gradCheckpoint

gradCheckpoint?: GradCheckpointTraining

`Optional` `Readonly`prefixShared

prefixShared?: PrefixSharedTraining

`Optional` `Readonly`segmented

segmented?: SegmentedTraining

Methods

`Optional`lmHead

lmHead?(): HeadQuant
The quantized LM head for fused/flash linear-CE heads.

Returns HeadQuant
- Defined in packages/inference/src/contracts/mlx/trainable.ts:97