Alex Spies afspies

Twitter thread: https://twitter.com/theshawwn/status/1456925974919004165
Hacker News thread: https://news.ycombinator.com/item?id=29128998

November 6, 2021

How does JAX allocate memory on a TPU?

jnp.device_put(1) is deceptively simple to write in JAX. But on a TPU, what actually happens? How does a tensor containing the value 1 actually get onto a TPU?

Turns out, the answer is "C++", and a lot of it.

They basically all suggest that apparent improvements to the state of the art in ML and related fields are often not real, or at least the result of factors other than what the authors claim.

The state of sparsity in deep neural networks

What is the state of neural network pruning?

On the State of the Art of Evaluation in Neural Language Models

Do Transformer Modifications Transfer Across Implementations and Applications?

UPDATE (March 2020, thanks @ic): I don't know the exact AMI version but yum install docker now works on the latest Amazon Linux 2. The instructions below may still be relevant depending on the vintage AMI you are using.

Amazon changed the install in Linux 2. One no-longer using 'yum' See: https://aws.amazon.com/amazon-linux-2/release-notes/

Docker CE Install

sudo amazon-linux-extras install docker
sudo service docker start

	# EDIT 10/04/2022 - This version was provided by @jayelm who fixed some bugs and made the function much more robust

	import os
	import subprocess
	import time

	def assign_free_gpus(threshold_vram_usage=1500, max_gpus=2, wait=False, sleep_time=10):
	"""
	Assigns free gpus to the current process via the CUDA_AVAILABLE_DEVICES env variable
	This function should be called after all imports,

	/* libtpujesus.c
	Copyright 2021 Shawn Presser

	Permission is hereby granted, free of charge, to any person obtaining
	a copy of this software and associated documentation files (the
	"Software"), to deal in the Software without restriction, including
	without limitation the rights to use, copy, modify, merge, publish,
	distribute, sublicense, and/or sell copies of the Software, and to
	permit persons to whom the Software is furnished to do so, subject to
	the following conditions:

	from imgaug import augmenters as iaa
	import imgaug as ia
	ia.seed(4)

	import tensorflow as tf
	tf.random.set_seed(666)

	aug = iaa.RandAugment(n=2, m=9)
	BATCH_SIZE = 224

	def downsample_2d(X, sz):
	"""
	Downsamples a stack of square images.

	Args:
	X: a stack of images (batch, channels, ny, ny).
	sz: the desired size of images.

	Returns:
	The downsampled images, a tensor of shape (batch, channel, sz, sz)

	"""Configuration library for experiments."""
	from typing import Dict, Any
	import logging
	import pprint
	import sys
	import argparse

	logger = logging.getLogger(__name__)

	parser = argparse.ArgumentParser(description=__doc__, fromfile_prefix_chars="@")