Forward pass of the mLSTM cell in recurrent form on a full sequence. The recurrent loop is implemented using a for loop. Args: m

Forward pass of the mLSTM cell in recurrent form on a full sequence. This function uses jax.lax.scan to loop over the sequence. Args:

Generated a denoised sample x_{t-1} given the trained model and noisy sample x_{t} :param x: noisy input action :param t:

Returns the names of available CLIP models

Returns list of utf-8 byte and a corresponding list of unicode strings. The reversible bpe codes work on unicode strings. This means you

Compile the provided function with CUDA graphs. Args: fn: The function to compile. Should take no arguments. warmups: The nu

Initialize optimizers and learning rate schedulers based on model configuration.

cosine schedule as proposed in https://openreview.net/forum?id=-NEXDKk8gZ

Create group names from columns in a DataFrame.

Calculates the exponential moving average over a vector. Will fail for large inputs. :param data: Input data :param alpha: scalar flo

This function indexes out a target dimension of a tensor in a structured way, by allowing a different value to be selected for each member of

Returns all actions from all trajectories, concatenated on dim 0 (time).

Constructs a continuous VP noise schedule.

Get a mlstm sequence kernel function by name. Naming convention: name = "<module_name>--<backend_name>" module_name: The name of th

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

Get the noise schedule for the sampling steps. Describes the distribution over the noise levels from sigma_min to sigma_max.

Get the noise schedule for the sampling steps. Describes the distribution over the noise levels from sigma_min to sigma_max.

Return set of symbol pairs in a word. Word is represented as tuple of symbols (symbols being variable-length strings).

Get all the parameter dictionaries for the benchmark.

Constructs an exponential noise schedule.

Constructs the noise schedule of Karras et al. (2022).

Constructs an linear noise schedule.

Constructs a continuous VP noise schedule.

Constructs a continuous VP noise schedule.

Validates the given kernel parameters or selects kernel params from heuristic. Either specify all kernel parameters or None. If None, the heuristi

Joins all dimensions between dimensions (@begin_axis, @end_axis) into a flat dimension, for all tensors in nested dictionary or list or tuple

Compute the loss for the denoising process. Args: state: The input state. action: The input action.

Execute the recurrent forward kernel for the C computation in the mLSTM chunkwise formulation. This function defines the grid and block size

Execute the recurrent forward kernel for the C computation in the mLSTM chunkwise formulation. This function defines the grid and block size

Computes the forward pass of the mLSTM chunkwise formulation. Args: matQ (jax.Array): The query matrix Q. Shape (B, NH, S, DHQK).

Pad input tensor or array @seq in the time dimension (dimension 1). Args: seq (np.ndarray or torch.Tensor): sequence to be padded

Convert a batch of (0th + Kth frame) images to their patched equivalents by naive reshaping.

Plot runtime results from a DataFrame. Convenience function to create a bar or line plot for runtime results. Also sets the matplotlib rc para

DPM-Solver 1( or DDIM sampler

A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022). SDE for S_churn!=0 and ODE otherwise 1. Last deno

Ancestral sampling with DPM-Solver inspired second-order steps. Ancestral sampling is based on the DDPM paper (https://arxiv.org/abs/2006.11

DPM-Solver-12 and 23 (adaptive step size). See https://arxiv.org/abs/2206.00927.

DPM-Solver-Fast (fixed step size). See https://arxiv.org/abs/2206.00927.

DPM-Solver++(2M).

DPM-Solver++(2M).

DPM-Solver++(2S) second-order steps.

Ancestral sampling combined with DPM-Solver++(2S) second-order steps.

DPM-Solver++ (stochastic).

Implements a variant of Algorithm 2 (Euler steps) from Karras et al. (2022). Stochastic sampler, which combines a first order ODE solver with

Ancestral sampling with Euler method steps. 1. compute dx_{i}/dt at the current timestep 2. get \sigma_{up} and \sigma_{down} from

Implements Algorithm 2 (Heun steps) from Karras et al. (2022). For S_churn =0 this is an ODE solver otherwise SDE Every update consists o

A linear multistep sampler. 1. compute x_{t-1} using the current noise level 2. compute dx/dt at x_{t-1} using the current noise level

Sets up the benchmark function to run.