import { GPTConfig } from '../models/config';
import { default as MemoryProfiler } from '../utilities/profile';
import { default as RoPECache } from './RoPECache';
import { Tensor, Variable } from '@tensorflow/tfjs-core';
import { default as WeightStore } from './WeightStore';
export interface ForwardAttributes {
    training: boolean;
    checkpointing?: boolean;
    mixedPrecision?: boolean;
    ropeCache?: RoPECache;
    outputEmbeddings?: boolean;
    embeddings?: {
        name: string;
        tensor: Tensor;
    }[];
    dropout?: number;
    layerDrop?: number;
}
export default abstract class BaseLayer<ATTR extends ForwardAttributes = ForwardAttributes, CONFIG extends GPTConfig = GPTConfig> {
    readonly parent?: BaseLayer;
    readonly config: CONFIG;
    weightStore: WeightStore;
    readonly children: BaseLayer[];
    private profiler?;
    private ownVariables;
    constructor(config: CONFIG, parent?: BaseLayer);
    getProfiler(): MemoryProfiler | undefined;
    setProfiler(profiler: MemoryProfiler | null): void;
    startMemory(): void;
    endMemory(label: string): void;
    addVariable(name: string, variable?: Variable): void;
    addChildVariable(name: string): void;
    get variables(): Variable[];
    get trainableVariables(): Variable[];
    getVariable(name: string): Tensor;
    hasVariable(name: string): boolean;
    setVariable(name: string, variable: Variable): void;
    dispose(): void;
    protected build(): void;
    abstract forward(attrs: ATTR, ...x: Tensor[]): Tensor | Tensor[];
    call(attrs: ATTR, ...x: Tensor[]): Tensor | Tensor[];
    callCheckpoint(attrs: ATTR, ...x: Tensor[]): Tensor;
    private checkpointingFn;
}
