norabelrose

from argparse import ArgumentParser
from itertools import pairwise
from pathlib import Path
from typing import Callable, Sized
import random

import pytorch_lightning as pl
import torch
import torch.nn.functional as F
import torchmetrics as tm

norabelrose

from dataclasses import dataclass, field

import torch
from torch import Tensor
from torch.nn.functional import (
    binary_cross_entropy_with_logits as bce_with_logits,
)
from torch.nn.functional import (
    cross_entropy,
)

norabelrose

from argparse import ArgumentParser
from typing import Any, Callable, Protocol, Sized, Type

import pytorch_lightning as pl
import torch
import torch.nn.functional as F
import torchmetrics as tm
import torchvision as tv
from concept_erasure import LeaceFitter, OracleFitter, QuadraticFitter
from pytorch_lightning.loggers import WandbLogger

norabelrose

from argparse import ArgumentParser
from itertools import pairwise
from typing import Callable, Sized

import pytorch_lightning as pl
import torch
import torch.nn.functional as F
import torchmetrics as tm
import torchvision as tv
from concept_erasure import LeaceFitter, OracleFitter, QuadraticFitter

norabelrose

from einops import rearrange
from torch import Tensor
import torch


def kronecker_decompose(
    A: Tensor, m: int, n: int, *, k: int = 1, niter: int = 10
) -> tuple[Tensor, Tensor]:
    """Frobenius-optimal decomposition of `A` into a sum of `k` Kronecker products.

norabelrose

from PIL import Image
from torch.utils.data import Dataset
import json
import os


class ImageNetKaggle(Dataset):
    def __init__(self, root, split, transform=None):
        self.samples = []
        self.targets = []

norabelrose

from torch import Tensor
import torch


def intersection_of_ranges(As: Tensor) -> Tensor:
    """Compute the intersection of the ranges of a batch of matrices.

    We use the formula from "Projectors on Intersections of Subspaces" by
    Ben-Israel (2015) <http://benisrael.net/ADI-BENISRAEL-AUG-29-13.pdf>.
    """

norabelrose

from dataclasses import dataclass
from typing import Callable, Iterator

from torch import LongTensor, Tensor
import torch


@dataclass(frozen=True)
class GroupedTensor:
    """A tensor split into groups along a given dimension.

norabelrose

import torch
from torch import Tensor, nn


class ConceptEraser(nn.Module):
    """Removes the subspace responsible for correlations between hiddens and labels."""

    mean_x: Tensor
    """Running mean of X."""

norabelrose

########################################################################################################
# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
########################################################################################################

import types, gc, os, time, re
import torch
from torch.nn import functional as F
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_tf32 = True