buttercutter

# Credit : gpt-4o

# Reference : [Protein Discovery with Discrete Walk-Jump Sampling](http://arxiv.org/abs/2306.12360)


import torch
import torch.nn as nn
import torch.optim as optim
import math

buttercutter

import ceviche
import matplotlib.pyplot as plt
import pandas as pd
from scipy.optimize import minimize
from autograd import grad
import autograd.numpy as anp


USE_ADAM = 0
USE_SCIPY = ~USE_ADAM

buttercutter

# [Nash Learning from Human Feedback](http://arxiv.org/abs/2312.00886)


import os
import math
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset

buttercutter

# Reference: [Half-Quadratic Quantization of Large Machine Learning Models](https://mobiusml.github.io/hqq_blog/)


import numpy as np

# Define the shrinkage function for soft-thresholding
def shrink(x, beta, p):
    return np.sign(x) * np.maximum(np.abs(x) - (np.abs(x)**(p-1))/beta, 0)

# Define the quantization and dequantization operators

buttercutter

# [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://arxiv.org/abs/2312.00752)

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torch.nn import functional as F
from einops import rearrange, repeat
from tqdm import tqdm

buttercutter

import torch
from torch import nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
from tqdm import tqdm

# Custom fast linear layer
class FastLinear(nn.Module):
    def __init__(self, in_features, out_features):

buttercutter

buttercutter

buttercutter

buttercutter