thomasahle

import torch
from functorch import vmap, grad
from torch.autograd import Function

sigmoid = torch.sigmoid
sigmoid_grad = vmap(vmap(grad(sigmoid)))

class TopK(Function):
    @staticmethod
    def forward(ctx, xs, k):

thomasahle

import torch.nn as nn

class TestModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.a = nn.Linear(10, 10)
        self.b = SubTestModule()

class SubTestModule(nn.Module):
    def __init__(self):

thomasahle

import numpy as np
from sklearn.cluster import KMeans
from sklearn.metrics.pairwise import euclidean_distances
import collections

def kl_means(X, k:int, l:int, policy:str):
    n, d = X.shape
    km = KMeans(k).fit(X)
    centers = km.cluster_centers_ / l
    labels = np.stack([km.labels_] * l).T

thomasahle

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import sklearn.metrics


def upper(xs, ys, convex=True):
    i = np.argsort(xs)
    ys = ys[i]
    xs = xs[i]

thomasahle

import numpy as np


def quartic(y, n1, n2, ip):
    """A solution to the equation
    n1 x2 + n2 (1 - x2) + 2 ip sqrt(x2 (1 - x2)) == y
    """
    assert n2 <= y <= n1
    d = np.sign(ip) * (ip**4 + ip**2 * (n1 - y) * (y - n2)) ** 0.5
    x2 = (2 * ip**2 + (n1 - n2) * (y - n2) - 2 * d) / (4 * ip**2 + (n1 - n2) ** 2)

thomasahle

import tqdm
import torch
import torch.nn as nn
import torch.optim as optim
import argparse
import torchdata.datapipes as dp
from torch.utils.data import DataLoader
from torch.nn import functional as F
import pytorch_lightning as pl
import random

thomasahle

\usepackage{graphicx}
\newcommand{\shortslash}{\raisebox{0.2ex}{\scalebox{0.65}{/}}}
\newcommand{\notdivides}{\!\mathrel{\backslash\kern-0.4em\shortslash}\!}

\newcommand{\notdividesTim}{\!\!\mathrel{\rotatebox[origin=c]{20}{$\nmid$}}\!\!}

\def\notdividesHeinrich{\mathpalette\notdiv\relax}
\let\divides=\backslash
\def\notdiv#1#2{\setbox0=\hbox{$#1\divides$}%
  \vcenter{\hbox to\wd0{\hss$\scriptscriptstyle/\hss$}}\kern-\wd0

thomasahle

def nonlocals():
    import inspect
     stack = inspect.stack()
     if len(stack) < 3: return {}
     f = stack[2][0]
     res = {}
     while f.f_back:
         res.update({k:v for k,v in f.f_locals.items() if k not in res})
         f = f.f_back
     return res

thomasahle

import numpy as np
import scipy.linalg as sla

theta_m = {3: 1.5e-2, 5: 5.4e-1, 7: 9.5e-1, 9: 2.1e0, 13: 5.4e0}

pade_coefficients = {
    3: [120, 60, 12, 1],
    5: [30240, 15120, 3360, 420, 30, 1],
    7: [17297280, 8648640, 1995840, 277200, 25200, 1512, 56, 1],
    9: [17643225600, 8821612800, 2075673600, 302702400, 30270240, 2162160, 110880, 3960, 90, 1],

thomasahle

import itertools

dp = [0] * 2**12
dp[0] = 1
for state in range(1, 2**12):
    for d1, d2 in itertools.product(range(6), repeat=2):
        o = 0
        s = d1 + d2 + 1
        if s < 12 and state & (1 << s):
            o = max(o, dp[state & ~(1 << s)])