Loads the schedulers state. Arguments: state_dict (dict): scheduler state. Should be an object returned from a ca

Returns the state of the scheduler as a :class:`dict`. It contains an entry for every variable in self.__dict__ which is not the opti

Performs a single optimization step. Arguments: closure (callable, optional): A closure that reevaluates the model

Ensure that 'number' can be 'divisor' divisible Reference from original tensorflow repo: https://github.com/tensorflow/models/blob/master

Computes covariance diagonal for Linear layer. Args: h: Input activations layer: Linear layer

Return correctly formatted lr/momentum for each param group.

Update the running average of parameters with a new value using a specified beta3 coefficient. This function is designed to update the runni

Return correctly formatted lr/momentum for each param group.

1x1 convolution

3x3 convolution with padding

Split the head vectors. This function splits the head vectors that have been concatenated (e.g. through the `merge_heads` function) i

Scales tensor values to range [0,1] using min-max normalization. Args: tensor: Input tensor epsilon: Small value to prevent

Delegates computation to layer-specific methods. Args: h: Input activations layer: PyTorch layer

Directly calls the instance to compute the diagonal and the full matrix of the gradient covariance matrix by delegating to layer-spec

A decoder layer. This module combines a Multi-headed Attention module and an MLP to create a layer of the transformer, with normaliza

Checks if the dimensions of a given parameter match the dimensions of its corresponding module's weights or bias. This function is cr

:param x: The input feature maps. (batch_size, in_c, h, w) :param kernel_size: the kernel size of the conv filter (tuple of two elements)

:param x: The input feature maps. (batch_size, in_c, h, w) :param kernel_size: the kernel size of the conv filter (tuple of two elements)

Computes the FIM for a given module's parameters. This method is called internally by the AdaFisher optimizer during the optimization

Updates a single parameter using AdaFisher optimization step. Args: hyperparameters: Dict containing 'lr', 'beta', and '

Updates a single parameter using AdaFisherW optimization step. Args: hyperparameters: Dict containing 'lr', 'beta', and

Performs a single optimization step for one parameter tensor, applying updates calculated from gradient and Fisher information.

Computes the precision@k for the specified values of k

Computes covariance diagonal for BatchNorm2d layer. Args: h: Input activations layer: BatchNorm2d layer

Computes the diagonal and the full matrix of the gradient covariance matrix for a BatchNorm2d layer. Parameters: - s (Tensor

Computes covariance diagonal for Conv2d layer. Args: h: Input (batch_size, in_channels, height, width) layer

Computes the diagonal and the full matrix of the gradient covariance matrix for a Conv2d layer. Parameters: - s (Tensor): Gr

r"""Densenet-121 model from `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_ Args: memory_efficie

Computes covariance diagonal for LayerNorm layer. Args: h: Input activations layer: LayerNorm layer

Computes the diagonal and the full matrix of the gradient covariance matrix for a LayerNorm layer. Parameters: - s (Tensor):

Computes the diagonal and the full matrix of the gradient covariance matrix for a Linear layer. Parameters: - s (Tensor): Gr

Loss function. This function computes the loss (negative log-likelihood). Parameters ---------- log_probas (`torch.F

r"""ResNet-101 model from `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_

r"""ResNet-50 model from `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_

Updates exponential moving average of parameters in-place. Args: new: New parameter values current: Current parameter values

Checks if the dimensions of a given parameter match the dimensions of its corresponding module's weights or bias. This function is cr

Extracts sliding window patches from input feature maps for convolution operations. Args: x: Input tensor (batch_size, in_channels,

Extract patches from input feature maps given a specified kernel size, stride, and padding. This function applies a sliding window approach

Computes Fisher Information Matrix with regularization for parameter updates. Args: module: Neural network layer

:param m: the layer :param classname: the class name of the layer :return: a matrix form of the gradient. it should be a [out

:param m: the layer :param classname: the class name of the layer :return: a matrix form of the gradient. it should be a [out

:param m: the layer :param p_grad_mat: the gradients in matrix form :return: a list of gradients w.r.t to the parameters in

:param m: the layer :param p_grad_mat: the gradients in matrix form :return: a list of gradients w.r.t to the parameters in

Initialize the weights