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 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

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

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 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)