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

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