norabelrose

import math
import numpy as np
from numpy.typing import ArrayLike, NDArray
from scipy.special import factorial2
from scipy.stats import norm


def relu_poly_ev(n: int, mu: ArrayLike, sigma: ArrayLike) -> NDArray:
    """
    Compute E[x^n * ReLU(x)] analytically where x ~ N(mu, sigma^2)

norabelrose

def x_gelu_expectation(mu, sigma):
    """Compute E[x * gelu(x)] for x ~ N(mu, sigma^2) analytically."""
    evCDF = norm.cdf(mu / np.sqrt(1 + sigma**2))
    evPDF = norm.pdf(mu / np.sqrt(1 + sigma**2)) / np.sqrt(1 + sigma**2)
    evZPDF = -mu*sigma/np.sqrt(1 + sigma**2)**3 * norm.pdf(mu / np.sqrt(1 + sigma**2))

    # linearity
    evXPDF = mu * evPDF + sigma * evZPDF

    # identity (first time)

norabelrose

import numpy as np
from scipy.stats import norm

def compute_E_xf(W1, W2, b1):
    """
    Computes the analytical expectation E[x f(x)^T] for a single hidden layer ReLU network.

    Parameters:
    - W1: numpy.ndarray, shape (k, n)
        Weight matrix of the first layer (W^{(1)}).

norabelrose

from argparse import ArgumentParser
from dataclasses import dataclass

import torch
import torchvision.transforms as T
from concept_erasure import QuadraticEditor, QuadraticFitter
from datasets import (
    ClassLabel, Dataset, DatasetDict, Features, Image, load_dataset
)
from einops import rearrange

norabelrose

from argparse import ArgumentParser
from pathlib import Path

from datasets import Dataset, load_dataset
from tqdm.auto import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

norabelrose

from argparse import ArgumentParser

from datasets import load_dataset
from peft import LoraConfig
from trl import DPOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments


if __name__ == "__main__":
    parser = ArgumentParser()

norabelrose

from itertools import pairwise
from typing import Literal

import pytorch_lightning as pl
import torch
import torchmetrics as tm
import torchvision as tv
from torch import nn
from torch.optim import RAdam
from torch.optim.lr_scheduler import CosineAnnealingLR

norabelrose

from itertools import (
    combinations_with_replacement as pyramid
)
from typing import Iterable
import math

from opt_einsum import get_symbol
from torch import Tensor
import torch

norabelrose

from itertools import product

import torch
import triton
import triton.language as tl


@triton.autotune(
    configs=[
        triton.Config({'BLOCK_N': n, 'BLOCK_D': d, 'GROUP_SIZE_D': 8}, num_stages=4, num_warps=4)

norabelrose

from typing import Optional
import torch


def get_all_the_cumulants(
    x: torch.Tensor, y: torch.Tensor, z: torch.Tensor, w: torch.Tensor, weights_in: Optional[torch.Tensor] = None
):
    if weights_in is not None:
        weights = weights_in
        weights = weights / weights.sum()