Skip to content

Instantly share code, notes, and snippets.

@kingoflolz
kingoflolz / shuffle_tfrecords.py
Created March 27, 2021 10:28
A quick script for shuffling tfrecord datasets
import tensorflow as tf
from tqdm import tqdm
index = open("data/openwebtext2_new_inputs.train.index").read().splitlines()
dataset = tf.data.Dataset.from_tensor_slices(index)
dataset = dataset.interleave(tf.data.TFRecordDataset, cycle_length=128, num_parallel_calls=tf.data.experimental.AUTOTUNE)
d = dataset.shuffle(10000).prefetch(100)
@kingoflolz
kingoflolz / simple_shard.py
Last active July 20, 2022 08:19
End to end example of model parallelism in Haiku + Jax, proving that gradients are allclose to the unsharded case
import os
# Run with a bunch of CPU devices.
def setUpModule():
global prev_xla_flags
prev_xla_flags = os.getenv("XLA_FLAGS")
flags_str = prev_xla_flags or ""
if "xla_force_host_platform_device_count" not in flags_str:
os.environ["XLA_FLAGS"] = (flags_str + " --xla_force_host_platform_device_count=4")
@kingoflolz
kingoflolz / main.py
Created March 13, 2021 10:52
jax-sharded-transformer
import os
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
# Run with a bunch of CPU devices.
def setUpModule():
global prev_xla_flags
@kingoflolz
kingoflolz / test.py
Last active August 26, 2020 13:27
multi tpu runtime
import os
import time
# os.environ["TF_CPP_MIN_VLOG_LEVEL"] = "2"
# TF 2.3.0, you can use kindiana/tf2_test venv to try it out
import tensorflow as tf
from cloud_tpu_client import Client
import logging
# from tensorflow.python.eager import context
/usr/bin/python3.8 /home/ben/src/fiducial_test/generate_images.py
266995
[examples/lftag.rs:27] idx = 0
[examples/lftag.rs:27] idx = 800
[examples/lftag.rs:27] idx = 2100
[examples/lftag.rs:27] idx = 700
[examples/lftag.rs:27] idx = 1600
[examples/lftag.rs:27] idx = 600
[examples/lftag.rs:27] idx = 1100
[examples/lftag.rs:27] idx = 900
[package]
name = "alarm"
version = "0.1.0"
authors = ["Ben Wang <[email protected]>"]
edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
mcp3xxx = "0.1.0-pre.1"
#include <cmath>
#include <cmath>
/**
Battery Tester 3001
*/
/* Includes ------------------------------------------------------------------*/
#include <cmsis_os.h>
import random
import subprocess
import numpy as np
import scipy.optimize
import os
import time
import copy
class Simulation:
def __init__(self):
self.commands = ["newdocument(0)", "showconsole()", 'mi_probdef(0, "millimeters", "planar", 1E-8)']
self.points = {}
self.materials = {}
self.circuits = 0
def run_sim(self):
self.commands += ['mi_saveas("Untitled.fem")', "mi_analyse()", "mi_loadsolution()"]
use std::collections::HashSet;
use super::*;
use std::collections::HashMap;
#[aoc_generator(day6)]
pub fn input_generator(input: &str) -> Vec<(usize, usize)> {
parse_shit(input.trim()).iter()
.map(|l| (l[0] + 1000, l[1] + 1000)).collect()
}