import { AttentionScores, CausalSelfAttentionConfig, KVCache } from './CausalSelfAttention';
import { MLPConfig } from './MLP';
import { RMSNormConfig } from './RMSNorm';
import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
import { Tensor } from '@tensorflow/tfjs-core';
import { GPTConfig } from '../models/config';
interface BlockAttributes extends ForwardAttributes {
    pastKV?: KVCache;
    seed?: number;
    attentionScores?: AttentionScores;
}
export type TransformerBlockConfig = MLPConfig & RMSNormConfig & CausalSelfAttentionConfig;
export default class Block extends BaseLayer<BlockAttributes> {
    private ln1;
    private attn;
    private ln2;
    private mlp;
    private index;
    skipped: boolean;
    private blockConfig;
    constructor(index: number, config: GPTConfig, blockConfig: TransformerBlockConfig, parent?: BaseLayer);
    private getMLPOutput;
    forward(attrs: BlockAttributes, x: Tensor): Tensor;
    dispose(): void;
}
export {};
