import { Tensor } from "./tensor";
export type StateDict = {
    [key: string]: Tensor;
};
/**
 * Base class for implementing a module, which is a reusable portion of a neural network.
 * Modules can contain other modules, allowing for a tree-like structure.
 *
 * To add a submodule to a module, assign it as a property of the module in the constructor.
 * Its name will be the name of the property.
 *
 * Modules can also contain `Parameter` objects, which are tensors that are automatically
 * updated by optimizers and are saved when calling `saveDict()`.
 * To add a parameter to a module, assign it as a property of the module in the constructor.
 *
 * Module subtypes implement a `forward` function in order to define the computation of the module.
 *
 * ```
 * class Model extends torch.nn.Module {
 *     conv1: torch.nn.Conv2d;
 *     conv2: torch.nn.Conv2d;
 *     constructor() {
 *         super();
 *         this.conv1 = torch.nn.Conv2d(1, 20, 5);
 *         this.conv2 = torch.nn.Conv2d(20, 20, 5);
 *     }
 *     forward(input: torch.Tensor): torch.Tensor {
 *         let output = this.conv1.forward(input);
 *         output = output.relu();
 *         output = this.conv2.forward(output);
 *         output = output.relu();
 *         return output;
 *     }
 * }
 * ```
 */
export declare class Module {
    private _children;
    /**
     * Returns the immediate children (submodules) of this module along with their names.
     */
    get namedChildren(): [string | number, Module][];
    /**
     * Returns the immediate children (submodules) of this module.
     */
    get children(): Module[];
    private _parameters;
    private get immediateParameters();
    private _buffers;
    private _nonPersistentBuffersSet;
    private _training;
    get training(): boolean;
    get [Symbol.toStringTag](): string;
    /**
     * Produces a multiline readable string of the module and its descendants.
     * @returns A string representation of the module and its descendants.
     */
    toString(): string;
    /**
     * Adds a child module of the current module.
     *
     * The module can be accessed as a property using the given name.
     * @param name name of the child module
     * @param module child module to be added
     */
    addModule(name: string | number, module: Module): void;
    /**
     * Returns this module and its descendants' along with their prefixed names.
     * @param memo is a set of modules used to avoid double counting.
     * @param prefix is prependended to the names of the modules.
     * @param removeDuplicate is a boolean indicating whether to remove duplicate modules.
     * @returns a generator of [prefixed name, module] pairs.
     */
    namedModules(memo?: Set<Module>, prefix?: string, removeDuplicate?: boolean): Generator<[string | number, Module]>;
    /**
     * Returns this module and its descendants.
     * @returns a generator of modules
     */
    modules(): Generator<Module>;
    private _named_members;
    /**
     * Gets this module and its descendants' (if `recurse = true`) parameters along with their prefixed names.
     * @param prefix is prependended to the names of the parameters
     * @param recurse whether to include submodule parameters
     * @param removeDuplicate whether to remove duplicate parameters
     * @returns a generator of [prefixed name, parameter] pairs
     */
    namedParameters(prefix?: string, recurse?: boolean, removeDuplicate?: boolean): Generator<[string, Parameter]>;
    parameters(): Generator<Parameter>;
    registerBuffer(name: string, tensor: Tensor | null, persistent?: boolean): void;
    namedBuffers(prefix?: string, recurse?: boolean, removeDuplicate?: boolean): Generator<[string, Tensor | null]>;
    buffers(): Generator<Tensor | null>;
    /**
     * Sets the module in training mode.
     * @param mode whether to set training mode (`true`) or evaluation mode (`false`).
     * @returns this module
     */
    train(mode?: boolean): ThisType<Module>;
    /**
     * Sets the module in evaluation mode.
     * @returns this module
     */
    eval(): ThisType<Module>;
    /**
     * Change if autograd should record operations on parameters in this module.
     * @param requiresGrad whether to enable gradient calculation for parameters in this module.
     * @returns this module
     */
    requiresGrad(requiresGrad?: boolean): ThisType<Module>;
    /**
     * Zeros out the gradients of all parameters. That can be accomplished either by
     * setting the `grad` property to `null` (`setToNull=true`) or
     * by filling the gradient tensor with zeros (`setToNull=false`).
     * @param setToNull whether to set gradients to `null` (`true`) or to zero tensors (`false`).
     * @returns this module
     */
    zeroGrad(setToNull?: boolean): ThisType<Module>;
    /**
     * Returns the state dictionary of the module and its descendants.
     * @param destination An optional state dictionary to update with the module's state.
     * @param prefix A string to prepend to the names of the state entries (default is an empty string).
     * @param keepVars A boolean flag, if true keeps the tensors attached to autograd (default is true).
     * @returns The updated state dictionary containing the module's state.
     */
    stateDict(destination?: StateDict, prefix?: string, keepVars?: boolean): StateDict;
    private _saveToStateDict;
    /**
     * Loads the state of the module and its descendants from the given state dictionary.
     * @param stateDict The state dictionary containing the state of the module to load.
     */
    loadStateDict(stateDict: StateDict): void;
    private _loadFromStateDict;
}
export declare class Parameter extends Tensor {
    constructor(data: Tensor, requiresGrad?: boolean);
}
/**
 * An abstraction for modules that accept child modules as arguments.
 */
export declare class Container extends Module {
}
/**
 * Hold submodules in a list.
 */
export declare class ModuleList extends Module {
    get length(): number;
    [index: number]: Module;
    [Symbol.iterator](): IterableIterator<Module>;
    constructor(modules?: Module[]);
    push(module: Module): void;
}
/**
 * A sequential container.
 *
 * Modules will be added to it in the order they are passed in the constructor.
 *
 * The `forward()` method of Sequential accepts any input and forwards it to the first module.
 * It then forwards the output of that module to the second module, and so on,
 * finally returning the output of the last module.
 */
export declare class Sequential extends Container {
    get length(): number;
    [index: number]: Module;
    [Symbol.iterator](): IterableIterator<Module>;
    constructor(modules: Module[]);
    forward(input: Tensor): Tensor;
}