Load and concatenate datasets from multiple directories. Args: base_dirs (list[str]): List of base directory paths containing the da

Load SAEs for multiple hookpoints on a single model and dataset.

Loads all .png or .jpg images from a given path Warnings: Expects all images to be of same dtype and shape. Args: path: relative path

Loads all .png or .jpg images from a given path Warnings: Expects all images to be of same dtype and shape. Args: path: relative path

Log disk I/O statistics using psutil and send them to wandb.

Concatenate a tensor across all processes.

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

Scatter & gather the hidden states across ranks.

dense1: shape (A, B) dense2: shape (B, N) at_indices: shape (A, K) out values: shape (A, K) calculates dense1 @ dense2 only for t

calculates sparse @ dense (i.e reducing along the uncollated dimension of sparse) dense must be contiguous along dim 0 (in other words, dense

calculates sparse.T @ dense (i.e reducing along the collated dimension of sparse) dense must be contiguous along dim 0 (in other words, dense

Simple hook that multiplies specific SAE feature activations by a scalar. Args: sae: The trained SAE model feature_id