<jupyter_start><jupyter_text><jupyter_code>import jax
import jax.numpy as jnp
from jax import value_and_grad
from jax import grad
from jax import vmap
from jax import jit
jax.device_count()
jax.devices()
import jax
key = jax.random.PRNGKey(0)
from numpy import random
import matplotlib.pyplot as plt
from tqdm.notebook import tqdm
#dir(jnp.zeros(5).device())
jnp.zeros(5).device().device_kind
#@title
# from https://github.com/google/jax/blob/main/examples/datasets.py

import array
import gzip
import os
from os import path
import struct
import urllib.request

import numpy as np


_DATA = "/tmp/jax_example_data/"


def _download(url, filename):
  """Download a url to a file in the JAX data temp directory."""
  if not path.exists(_DATA):
    os.makedirs(_DATA)
  out_file = path.join(_DATA, filename)
  if not path.isfile(out_file):
    urllib.request.urlretrieve(url, out_file)
    print("downloaded {} to {}".format(url, _DATA))


def _partial_flatten(x):
  """Flatten all but the first dimension of an ndarray."""
  return np.reshape(x, (x.shape[0], -1))


def _one_hot(x, k, dtype=np.float32):
  """Create a one-hot encoding of x of size k."""
  return np.array(x[:, None] == np.arange(k), dtype)


def mnist_raw():
  """Download and parse the raw MNIST dataset."""
  # CVDF mirror of http://yann.lecun.com/exdb/mnist/
  base_url = "https://storage.googleapis.com/cvdf-datasets/mnist/"

  def parse_labels(filename):
    with gzip.open(filename, "rb") as fh:
      _ = struct.unpack(">II", fh.read(8))
      return np.array(array.array("B", fh.read()), dtype=np.uint8)

  def parse_images(filename):
    with gzip.open(filename, "rb") as fh:
      _, num_data, rows, cols = struct.unpack(">IIII", fh.read(16))
      return np.array(array.array("B", fh.read()),
                      dtype=np.uint8).reshape(num_data, rows, cols)

  for filename in ["train-images-idx3-ubyte.gz", "train-labels-idx1-ubyte.gz",
                   "t10k-images-idx3-ubyte.gz", "t10k-labels-idx1-ubyte.gz"]:
    _download(base_url + filename, filename)

  train_images = parse_images(path.join(_DATA, "train-images-idx3-ubyte.gz"))
  train_labels = parse_labels(path.join(_DATA, "train-labels-idx1-ubyte.gz"))
  test_images = parse_images(path.join(_DATA, "t10k-images-idx3-ubyte.gz"))
  test_labels = parse_labels(path.join(_DATA, "t10k-labels-idx1-ubyte.gz"))

  return train_images, train_labels, test_images, test_labels


def mnist(permute_train=False):
  """Download, parse and process MNIST data to unit scale and one-hot labels."""
  train_images, train_labels, test_images, test_labels = mnist_raw()

  train_images = _partial_flatten(train_images) / np.float32(255.)
  test_images = _partial_flatten(test_images) / np.float32(255.)
  train_labels = _one_hot(train_labels, 10)
  test_labels = _one_hot(test_labels, 10)

  if permute_train:
    perm = np.random.RandomState(0).permutation(train_images.shape[0])
    train_images = train_images[perm]
    train_labels = train_labels[perm]

  return {'train_img': train_images, 'train_lab': train_labels, 'test_img': test_images, 'test_lab': test_labels}
mnist_dat = mnist()
mnist_dat['train_img'][0].shape
import matplotlib.pyplot as plt
def show_img(img):
  plt.imshow(img.reshape(28,28))
  plt.colorbar()
demo_idx = 33
demo_img = mnist_dat['train_img'][demo_idx]
show_img(demo_img)
print(jnp.argmax(mnist_dat['train_lab'][demo_idx]))
#jnp.argmax(model.forward_single(demo_img))
def model_random_predict(img):
  return int(random.random()*10)
model_random_predict('test')
def evaluate(pred_func, max=2000, progress_func=lambda x: x):
  correct = 0
  incorrect = 0
  count = 0

  for train_img, train_lab in progress_func(zip(mnist_dat['test_img'][:max], mnist_dat['test_lab'][:max])):
    pred = pred_func(train_img)
    if jnp.argmax(pred) == jnp.argmax(train_lab):
      correct += 1
    else:
      incorrect += 1
    count += 1
    final_stats = {'correct': correct, 'total': count, 'accuracy': correct/count}
    #tqdm.write(str(final_stats))

  return final_stats
evaluate(model_random_predict)
class NNModel:

  def __init__(self, input_size, hidden_size, out_size, ce_loss=False, use_mom=False):

    self.input_size = input_size
    self.hidden_size = hidden_size
    self.out_size = out_size

    self.cross_entropy_loss = ce_loss
    self.momentum = use_mom

    # use normed initialization?
    self.weights_l1 = jnp.array(random.normal(size=(input_size, hidden_size))) / jnp.sqrt(hidden_size)
    self.w1_m = jnp.zeros_like(self.weights_l1)
    self.weights_l2 = jnp.array(random.normal(size=(hidden_size, out_size))) / jnp.sqrt(out_size)
    self.w2_m = jnp.zeros_like(self.weights_l2)
    self.bias_l1 = jnp.array(random.normal(size=(hidden_size,)))
    self.b1_m = jnp.zeros_like(self.bias_l1)
    self.bias_l2 = jnp.array(random.normal(size=(out_size,)))
    self.b2_m = jnp.zeros_like(self.bias_l2)

    @jit
    def forward_single(weights_l1, weights_l2, bias_l1, bias_l2, img):
      # matrix mult, bias, relu activation
      hidden = jnp.maximum(img @ weights_l1 + bias_l1, 0.0)
      out_raw = hidden @ weights_l2 + bias_l2
      # softmax to probability distribution
      out_exp = jnp.e ** out_raw
      out_probs = out_exp / out_exp.sum()
      return out_probs

    @jit
    def loss(w1, w2, b1, b2, img, lab):
      preds = forward_single(w1, w2, b1, b2, img)
      #print(preds)
      #print(lab)
      if self.cross_entropy_loss:
        idx = jnp.argmax(lab)
        return -jnp.log(preds[idx])
      else:
        diffs = (preds-lab) ** 2
        return diffs.mean()

    self.forward_single = lambda img: forward_single(
        self.weights_l1, self.weights_l2, 
        self.bias_l1, self.bias_l2, img)
    
    self.grad_f = value_and_grad(loss, argnums=[0,1,2,3])
    self.grad_func = lambda img, lab: self.grad_f(
        self.weights_l1, 
        self.weights_l2, 
        self.bias_l1, 
        self.bias_l2, 
        img, lab)

    self.batched_loss = vmap(loss, in_axes=(None, None, None, None, 0, 0))

    self._grad_batched = jit(value_and_grad(
        lambda w1, w2, b1, b2, imgs, labs: self.batched_loss(w1, w2, b1, b2, imgs, labs).sum(),
        argnums=[0,1,2,3]
    ))

    self.grad_batched = lambda imgs, labs: self._grad_batched(
            self.weights_l1, self.weights_l2, 
            self.bias_l1, self.bias_l2, imgs, labs)

  def update(self, imgs, labs, lr_a, lr_b):
    #loss, grads = self.grad_func(imgs, labs)
    loss, grads = self.grad_batched(imgs, labs)
    w1_d, w2_d, b1_d, b2_d = grads
    if self.momentum:
      # update momentum
      self.w1_m = lr_b * self.w1_m + w1_d * (1-lr_b)
      self.w2_m = lr_b * self.w2_m + w2_d * (1-lr_b)
      self.b1_m = lr_b * self.b1_m + b1_d * (1-lr_b)
      self.b2_m = lr_b * self.b2_m + b2_d * (1-lr_b)
      # update parameters
      self.weights_l1 -= self.w1_m * lr_a
      self.weights_l2 -= self.w2_m * lr_a
      self.bias_l1 -= self.b1_m * lr_a
      self.bias_l2 -= self.b2_m * lr_a
    else:
      # update parameters
      self.weights_l1 -= w1_d * lr_a #self.w1_m * lr_b
      self.weights_l2 -= w2_d * lr_a #self.w2_m * lr_b
      self.bias_l1 -= b1_d * lr_a #self.b1_m * lr_b
      self.bias_l2 -= b2_d * lr_a #self.b2_m * lr_b
    return loss

def create_fresh_mnist_model(hidden_size=256, use_momentum=False, use_ce_loss=False):
  return NNModel(
      mnist_dat['train_img'][0].shape[0], hidden_size, mnist_dat['train_lab'][0].shape[0],
      ce_loss=use_ce_loss, use_mom=use_momentum)
model = create_fresh_mnist_model()
loss_single, grd_single = model.grad_func(mnist_dat['train_img'][demo_idx], mnist_dat['train_lab'][demo_idx])
for x in grd_single:
  print(x.shape)
  print(x.std())
loss_single
loss_b, grd_b = model.grad_batched(mnist_dat['train_img'][demo_idx:demo_idx + 1], mnist_dat['train_lab'][demo_idx:demo_idx + 1])
loss_b
for grds, grdb in zip(grd_single, grd_b):
  print((grds-grdb).sum())
len(mnist_dat['train_img'])
def test_train_performance(train_its, test_interval, test_its, batch_size=16, hidden_size=256, use_momentum=False, use_ce_loss=False, alpha=0.005, beta=0.9):
  mod = create_fresh_mnist_model(hidden_size=hidden_size, use_momentum=use_momentum, use_ce_loss=use_ce_loss)
  accs = []
  for i in tqdm(range(train_its)):
    lb, ub = i*batch_size, (i+1)*batch_size
    #loss = mod.update(mnist_dat['train_img'][i], mnist_dat['train_lab'][i], alpha, beta)
    loss = mod.update(mnist_dat['train_img'][lb:ub], mnist_dat['train_lab'][lb:ub], alpha, beta)
    if i % test_interval == 0:
      accs.append(evaluate(mod.forward_single, max=test_its))
  return {
      'stats': accs, 
      'train_its': train_its, 'test_interval': test_interval, 
      'test_its': test_its, 'hidden_size': hidden_size, 
      'use_momentum': use_momentum, 'use_ce_loss': use_ce_loss, 
      'alpha': alpha, 'beta': beta }
experiments = {}
#experiments['lr_0.005'] = test_train_performance(2500, 500, 200, alpha=0.005)
#experiments['lr_0.015'] = test_train_performance(2500, 500, 200, alpha=0.015)
experiments['lr_0.05_b1'] = test_train_performance(2500, 500, 200, alpha=0.05, batch_size=1)
experiments['lr_0.05'] = test_train_performance(2500, 500, 200, alpha=0.05, batch_size=16)
#experiments['lr_0.1'] = test_train_performance(2500, 500, 200, alpha=0.1)
experiments['lr_0.05_momentum'] = test_train_performance(2500, 500, 200, use_momentum=True, alpha=0.05, beta=0.9, batch_size=1)
experiments['lr_0.005_ce_loss'] = test_train_performance(2500, 500, 200, use_ce_loss=True, alpha=0.005)
experiments['lr_0.005_ce_loss_momentum'] = test_train_performance(2500, 500, 200, use_ce_loss=True, use_momentum=True, alpha=0.005, beta=0.9)
experiments['lr_0.005_ce_loss_momentum'] = test_train_performance(2500, 500, 200, use_ce_loss=True, use_momentum=True, alpha=0.005, beta=0.9)
experiments['lr_0.002_ce_loss_momentum'] = test_train_performance(2500, 500, 200, use_ce_loss=True, use_momentum=True, alpha=0.002, beta=0.9)
experiments['lr_0.002_momentum'] = test_train_performance(2500, 500, 200, use_momentum=True, alpha=0.002, beta=0.9)
for name, data in experiments.items():
  nums = [s['accuracy'] for s in data['stats']]
  plt.plot(list(range(len(nums))), nums, label=name)
plt.legend()
experiments
def funcy(x):
  return jnp.sum(jnp.cos(x)+1.5)**2

def mag(f):
  return lambda x: f(x).mean()

fg = grad(funcy)
bad = grad(grad(funcy))
mega = grad(mag(grad(mag(grad(funcy)))))
inp = jnp.arange(0,10, 0.1, dtype=jnp.float32)
fg(inp)
mega(inp)
def proc_vec(va, vb):
  return va.mean() + vb.mean()

b_vec = vmap(value_and_grad(proc_vec))

ina = jax.random.normal(key, (16, 1784))
inb = jax.random.normal(key, (16, 64))

b_vec(ina, inb)
#proc_vec(jnp.arange(10, dtype=jnp.float32))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14244240494474622, 'fraction_numerical': 0.038677654991571456, 'mean_word_length': 4.307654075546719, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '195016', 'license': 'MIT', 'n_tokens_mistral': 4496, 'n_tokens_neox': 4229, 'n_words': 785, 'path': 'nn_from_scratchV3.ipynb', 'repo_name': 'PWhiddy/jax-experiments', 'size': '10678'}