KeAWang

import torch
from torch.utils.flop_counter import FlopCounterMode
from triton.testing import do_bench

def get_flops_achieved(f):
    flop_counter = FlopCounterMode(display=False)
    with flop_counter:
        f()
    total_flops = flop_counter.get_total_flops()
    ms_per_iter = do_bench(f)

KeAWang

# %%
import torch
import numpy as np
def make_adding_dataset(num_seqs, seq_len, num_terms=2, seed=43141):
    assert 0 <= num_terms <= seq_len
    rng = np.random.default_rng(seed=seed)
    numbers = rng.uniform(0, 1, (num_seqs, seq_len))  # B x T
    mask = np.zeros_like(numbers)  # B x T
    non_zero = np.stack([rng.choice(seq_len, num_terms, replace=True) for _ in range(num_seqs)])  # B x 2
    mask[np.arange(num_seqs)[:, None], non_zero] = 1  # mask[i, non_zero[i, j]] 

KeAWang

import torch

class NanWrapper(torch.nn.Module):
    """Wrapper module around a torch Module that handles incoming nans"""
    def __init__(self, module):
        super().__init__()
        self.module = module

    def forward(self, x):
        """ Masks the entire last dimension (usually the feature/channel dimension) if any element is NaN. """

KeAWang

import torch


def count_params(model: torch.nn.Module):
  """count number trainable parameters in a pytorch model"""
  total_params = sum(torch.numel(x) for x in model.parameters())
  return total_params

KeAWang

import torch


def constrain(x, min, max, temperature:float=1.):
    return (max - min) * torch.sigmoid(x / temperature) + min

def unconstrain(y, min, max, temperature:float=1, EPS:float=1e-8):
    assert torch.all(y >= min) and torch.all(y <= max)
    # ensure both numerator and denominator are positive
    numerator = y - min 

KeAWang

import numpy as np
tir_palette = {
    "very_low": "#A61D2A",
    "low": "#EE1D23", 
    "in_range": "#26B257",
    "high": "#FAAB1A",
    "very_high": "#F47D21"
}
def color_bg(bgs):
    bins = [0, 54, 70, 180, 250, 1000]

KeAWang

import pandas as pd
from typing import Optional, List


def array_to_dataframe(array, axis_names: Optional[List[str]]=None):
    """Based on https://stackoverflow.com/questions/35525028/how-to-transform-a-3d-arrays-into-a-dataframe-in-python"""
    if axis_names is None:
        axis_names = list(range(array.ndim))
    index = pd.MultiIndex.from_product([range(s) for s in array.shape], names=names)
    df = pd.DataFrame({"array": array.flatten()}, index=index)["array"]

KeAWang

# %%
import jax.numpy as jnp
import jax
import equinox as eqx
from typing import Union, Any
from abc import ABC, abstractmethod

MaybeParameterizedArray = Union[jax.Array, "ParameterizedArray"]

KeAWang

from collections import namedtuple


class Data(namedtuple("Data", ("x", "y"))):
     def to(self, device, non_blocking=False):
         x = self.x.to(device, non_blocking=non_blocking)
         y = self.y.to(device, non_blocking=non_blocking)
         return Data(x, y)

     def contiguous(self):

KeAWang

import os
import pickle
from pathlib import Path
from wandb.errors import CommError

import wandb


def get_history(user="", project="", query={}, **kwargs):
    api = wandb.Api()