Flops for a single sequence.

Flops for a single sequence.

Counts the number of flops in the forward pass of an sLSTM cell.

DeepMind method for forwad pass FLOPs counting of decoder-only Transformer See Chinchilla paper or this blog post: https://www.adamcasson.com/

Create a bar plot for runtime results. Example usage with rc_context_wrapper: ```python fig = rc_context_wrapper( func=create_run

Create a line plot for runtime results. Simliar to `create_runtime_bar_plot`, but creates a line plot instead of a bar plot. Args: da

Crops images at the locations specified by @crop_indices. Crops will be taken across all channels. Args: images (torch.Tensor):

Denoise the next sequence of actions

Main Method to generate actions conditioned on the batch of state inputs :param state: the current state observation to conditon the

Denoise the next sequence of actions

Denormalize a tensor image with mean and standard deviation. Args: tensor (Tensor): Normalized image tensor. means (list): Me

Computes the score matching loss given the perceptual embedding, latent goal, and desired actions.

Computes the score matching loss given the perceptual embedding, latent goal, and desired actions.

Expand all tensors in nested dictionary or list or tuple at a single dimension @dim by @size. Args: x (dict or list or tuple): a

Expand a tensor at a single dimension @dim by @size Args: x (torch.Tensor): input tensor size (int): size to expand

Feed forward layer with given activation function

Flatten all tensors in nested dictionary or list or tuple, from @begin_axis onwards. Args: x (dict or list or tuple): a possibly nes

Flatten a nested dict or list to a list. For example, given a dict { a: 1 b: { c: 2 } c: 3

Flatten a tensor in all dimensions from @begin_axis onwards. Args: x (torch.Tensor): tensor to flatten begin_axis (int): whi

When using network mount from nautilus, format path

Forward pass of the mLSTM backend. Depending on the configured mode, this method will call the appropriate kernel function. Args:

Samples N random crops for each input in the batch, and then reshapes inputs to [B * N, ...].

Apply @gather_along_dim_with_dim_single to all tensors in a nested dictionary or list or tuple. Args: x (dict or list or tuple):

Get a list of available mlstm sequence kernels. These kernels process a sequence in the parallel or chunkwise parallel mode of the mLSTM.

Returns the available mlstm step kernels. These kernels can be used to compute a single time step of the mLSTM, i.e. for generation.

Get a mlstm sequence kernel function by name. Naming convention: name = "<backend_name>" backend_name: The name of the kernel funct

Returns the coefficient of the j-th derivative of the i-th step of a linear multistep method.

Method to load pretrained weights for the entire agent

Run main training/testing pipeline.

Generate a sample density function based on the desired type for training the model We mostly use log-logistic as it has no additiona

Generate a sample density function based on the desired type for training the model We mostly use log-logistic as it has no additiona

Perform per-context random masking by shuffling :: uses argsort random noise to identify masked patches.

Benchmark the runtime of the provided function. By default, return the mean. Copy with minor adaptations. of the original function from trit

This function defines the grid and block sizes for the kernel launch and calls the kernel. chunk parallel size: siz_b_LKV chunk loop si

Computes only the deltaK gradients for the backward pass. The other gradients are computed in the other (kernel) function. This function

This function defines the grid and block sizes for the kernel launch and calls the kernel. chunk parallel size: siz_b_LQ chunk loop siz

Computes only the deltaK gradients for the backward pass. The other gradients are computed in the other (kernel) function. This function

Computes the gradients for the query, key and value matrices. Args: matQ: Tensor containing the query vectors. Shape (B, NH, S, DHQK

This function defines the grid and block sizes for the kernel launch and calls the kernel. chunk parallel size: siz_b_LKV chunk loop si

Computes only the deltaK gradients for the backward pass. The other gradients are computed in the other (kernel) function. This function

This function defines the grid and block sizes for the kernel launch and calls the kernel.

Execute the parallel forward kernel for the H computation in the mLSTM chunkwise formulation. This function defines the grid and block sizes

This function computes the output of the mLSTM chunkwise formulation. It is the second part of the chunkwise mLSTM forward pass and combines the