Save the SAEs to disk.

Select the top-k latents.

Encode the input and select the top-k latents.

Run the pipeline with hooks at specified positions. Returns the final output. Args: *args: Arguments to pass to

Run the pipeline with caching at specified positions. This method allows you to cache the intermediate inputs and/or outputs of the

Locate the block at the specified position in the pipeline.

Calculates the statistics used by FID Args: images: torch.tensor, shape: (N, 3, H, W), dtype: torch.float32 in range 0 - 1 batch_s

Calculate FID between images1 and images2 Args: images1: np.array, shape: (N, H, W, 3), dtype: np.float32 between 0-1 or np.uint8

Returns the threshold for the given percentile. Args: scores (torch.Tensor): 1D tensor of unnormalized scores, shape [num_features].

Resizes and shifts the dynamic range of image to 0-1 Args: images: np.array, shape: (N, H, W, 3), dtype: float32 between 0-1 or np.uint8

Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles custom timesteps. Any kwargs

Transfers elements to cuda if GPU is available Args: elements: torch.tensor or torch.nn.module -- Returns: elemen

Consolidate sharded datasets into a single directory without rewriting data. Each of the shards must be of the same format, aka the full data

Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.

Locate the block at the specified position in the pipeline.

Make causal mask used for bi-directional self-attention.

bisection method to find the root for the projection operation of $u$

Numpy implementation of the Frechet Distance. The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1) and X_2 ~ N(mu_2, C_2

Prepare a plan for distributing modules across ranks.

Compute the geometric median `points`. Used for initializing decoder bias.

Calculates activations for last pool layer for all iamges -- Images: torch.array shape: (N, 3, 299, 299), dtype: torch.float32