r"""Returns the mean of the scalars that were accumulated for the given statistic between the last two calls to `update()`, or NaN if

Wrapper function: transfrom velocity prediction model to score Args: velocity: [batch_dim, ...] shaped tensor; velocity model output

Write text to stdout (and a file) and optionally flush.

r"""Copies current values of the internal counters to the user-visible state and resets them for the next round. If `keep_previous=Tr

r"""Returns the raw moments that were accumulated for the given statistic between the last two calls to `update()`, or zero if no scal

Traverses the object name and returns the last (rightmost) python object.

Take the mean over all non-batch dimensions.

Compute pairwise distances between two batches of feature vectors.

Flush, close possible files, and remove stdout/stderr mirroring.

Compute image features for downstream evals. :param batches: a iterator over NHWC numpy arrays in [0, 255]. :return: a tuple

Evaluate precision and recall efficiently. :param features_1: [N1 x D] feature vectors for reference batch. :param radii_1:

Flush written text to both stdout and a file, if open.

Compute the Frechet distance between two sets of statistics.

embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)

Searches for the underlying module behind the name to some python object. Returns the module and the object name (original name with module part r

r"""Test whether the given object or class is persistent, i.e., whether it will save its source code when pickled.

r"""Returns the names of all statistics broadcasted so far that match the regular expression specified at construction time.

Evaluate pairwise distances between two batches of feature vectors.

Step the EMA model towards the current model.

r"""Check that the given object is pickleable, raising an exception if it is not. This function is expected to be considerably more efficient

Randomly drop a fraction of tokens (except we ensure at least one token kept). x: (B, T, C) drop_ratio: fraction of tokens t

r"""Query the source code of a given Python module.

x_sparse: (B, T_keep, C) ids_keep: (B, T_keep) T_full: full sequence length T Returns: x_pad: (B, T_full,

Copied from: https://github.com/numpy/numpy/blob/fb215c76967739268de71aa4bda55dd1b062bc2e/numpy/lib/format.py#L788-L886 Read from file-like

f_dense: (B, T, C) encoder output ft g_full: (B, T, C) padded sparse output g_pad Returns fused h: (B, T, C)

r"""Get or create a Python module for the given source code.

r"""Synchronize the global cumulative counters across devices and processes. Called internally by `Collector.update()`.

Finds the python object with the given name and calls it as a function.

Center crop and resize image.

Center cropping implementation from ADM. https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diff

Center crop wide image.

Dataset processing tool for dataset image data conversion and VAE encode/decode preprocessing.

Create a logger that writes to a log file and stdout.

Builds a single .npz file from a folder of .png samples.

Downloads a pre-trained checkpoint from Hugging Face Hub and returns the local file path.

Function to reshape time t to broadcastable dimension of x Args: t: [batch_dim,], time vector x: [batch_dim,...], data point

grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/

Finds the python object with the given name.