<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 <a href="https
Step1:  Import libraries
Step4:  Define helper functions for processing data and defining a vanilla momentum optimizer
Step5:  Function Space Linearization
Create MNIST data pipeline using TensorFlow Datasets.
Step6:  Setup some experiment parameters.
Step7:  Create a Fully-Connected Network.
Step8:  Construct the NTK.
Step9:  Now that we have the NTK and a network we can compare against a number of different dynamics. Remember to reinitialize the network and NTK if you want to try a different dynamics.
Gradient Descent, MSE Loss
Create a optimizer and initialize it.
Step10:  Create an MSE loss and a gradient.
Step11:  Create an MSE predictor and compute the function space values of the network at initialization.
Step12:  Train the network.
Step13:  Gradient Descent, Cross Entropy Loss
Create a optimizer and initialize it.
Step14:  Create an Cross Entropy loss and a gradient.
Step15:  Create a Gradient Descent predictor and compute the function space values of the network at initialization.
Step16:  Train the network.
Step17:  Momentum, Cross Entropy Loss
Create a optimizer and initialize it.
Step18:  Create a Cross Entropy loss and a gradient.
Step19:  Create a momentum predictor and initialize it.
Step20:  Train the network.
<ASSISTANT_TASK:>
Python Code:
!pip install --upgrade pip
!pip install -q tensorflow-datasets
!pip install --upgrade jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_releases.html
!pip install -q git+https://www.github.com/google/neural-tangents
Explanation: <a href="https://colab.research.google.com/github/google/neural-tangents/blob/main/notebooks/function_space_linearization.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
Copyright 2019 Google LLC.
Licensed under the Apache License, Version 2.0 (the "License");
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
Import & Utils
Install JAX, Tensorflow Datasets, and Neural Tangents.
The first line specifies the version of jaxlib that we would like to import. Note, that "cp36" species the version of python (version 3.6) used by JAX. Make sure your colab kernel matches this version.
End of explanation
from jax import jit
from jax import grad
from jax import random
import jax.numpy as np
from jax.example_libraries.stax import logsoftmax
from jax.example_libraries import optimizers
import tensorflow_datasets as tfds
import neural_tangents as nt
from neural_tangents import stax
Explanation: Import libraries
End of explanation
def process_data(data_chunk):
  Flatten the images and one-hot encode the labels.
  image, label = data_chunk['image'], data_chunk['label']
  samples = image.shape[0]
  image = np.array(np.reshape(image, (samples, -1)), dtype=np.float32)
  image = (image - np.mean(image)) / np.std(image)
  label = np.eye(10)[label]
  return {'image': image, 'label': label}
@optimizers.optimizer
def momentum(learning_rate, momentum=0.9):
  A standard momentum optimizer for testing.
  Different from `jax.example_libraries.optimizers.momentum` (Nesterov).
  
  learning_rate = optimizers.make_schedule(learning_rate)
  def init_fn(x0):
    v0 = np.zeros_like(x0)
    return x0, v0
  def update_fn(i, g, state):
    x, velocity = state
    velocity = momentum * velocity + g
    x = x - learning_rate(i) * velocity
    return x, velocity
  def get_params(state):
    x, _ = state
    return x
  return init_fn, update_fn, get_params
Explanation: Define helper functions for processing data and defining a vanilla momentum optimizer
End of explanation
dataset_size = 64
ds_train, ds_test = tfds.as_numpy(
    tfds.load('mnist:3.*.*', split=['train[:%d]' % dataset_size,
                                    'test[:%d]' % dataset_size],
              batch_size=-1)
)
train = process_data(ds_train)
test = process_data(ds_test)
Explanation: Function Space Linearization
Create MNIST data pipeline using TensorFlow Datasets.
End of explanation
learning_rate = 1e0
training_steps = np.arange(1000)
print_every = 100.0
Explanation: Setup some experiment parameters.
End of explanation
init_fn, f, _ = stax.serial(
    stax.Dense(512, 1., 0.05),
    stax.Erf(),
    stax.Dense(10, 1., 0.05))
key = random.PRNGKey(0)
_, params = init_fn(key, (-1, 784))
Explanation: Create a Fully-Connected Network.
End of explanation
ntk = nt.batch(nt.empirical_ntk_fn(f, vmap_axes=0),
               batch_size=64, device_count=0)
g_dd = ntk(train['image'], None, params)
g_td = ntk(test['image'], train['image'], params)
Explanation: Construct the NTK.
End of explanation
opt_init, opt_apply, get_params = optimizers.sgd(learning_rate)
state = opt_init(params)
Explanation: Now that we have the NTK and a network we can compare against a number of different dynamics. Remember to reinitialize the network and NTK if you want to try a different dynamics.
Gradient Descent, MSE Loss
Create a optimizer and initialize it.
End of explanation
loss = lambda fx, y_hat: 0.5 * np.mean((fx - y_hat) ** 2)
grad_loss = jit(grad(lambda params, x, y: loss(f(params, x), y)))
Explanation: Create an MSE loss and a gradient.
End of explanation
predictor = nt.predict.gradient_descent_mse(g_dd, train['label'],
                                            learning_rate=learning_rate)
fx_train = f(params, train['image'])
Explanation: Create an MSE predictor and compute the function space values of the network at initialization.
End of explanation
print ('Time\tLoss\tLinear Loss')
X, Y = train['image'], train['label']
predictions = predictor(training_steps, fx_train)
for i in training_steps:
  params = get_params(state)
  state = opt_apply(i, grad_loss(params, X, Y), state)
  if i % print_every == 0:
    exact_loss = loss(f(params, X), Y)
    linear_loss = loss(predictions[i], Y)
    print('{}\t{:.4f}\t{:.4f}'.format(i, exact_loss, linear_loss))
Explanation: Train the network.
End of explanation
opt_init, opt_apply, get_params = optimizers.sgd(learning_rate)
state = opt_init(params)
Explanation: Gradient Descent, Cross Entropy Loss
Create a optimizer and initialize it.
End of explanation
loss = lambda fx, y_hat: -np.mean(logsoftmax(fx) * y_hat)
grad_loss = jit(grad(lambda params, x, y: loss(f(params, x), y)))
Explanation: Create an Cross Entropy loss and a gradient.
End of explanation
predictor = nt.predict.gradient_descent(loss, g_dd, train['label'], learning_rate=learning_rate)
fx_train = f(params, train['image'])
Explanation: Create a Gradient Descent predictor and compute the function space values of the network at initialization.
End of explanation
print ('Time\tLoss\tLinear Loss')
X, Y = train['image'], train['label']
predictions = predictor(training_steps, fx_train)
for i in training_steps:
  params = get_params(state)
  state = opt_apply(i, grad_loss(params, X, Y), state)
  if i % print_every == 0:
    t = i * learning_rate
    exact_loss = loss(f(params, X), Y)
    linear_loss = loss(predictions[i], Y)
    print('{:.0f}\t{:.4f}\t{:.4f}'.format(i, exact_loss, linear_loss))
Explanation: Train the network.
End of explanation
mass = 0.9
opt_init, opt_apply, get_params = momentum(learning_rate, mass)
state = opt_init(params)
Explanation: Momentum, Cross Entropy Loss
Create a optimizer and initialize it.
End of explanation
loss = lambda fx, y_hat: -np.mean(logsoftmax(fx) * y_hat)
grad_loss = jit(grad(lambda params, x, y: loss(f(params, x), y)))
Explanation: Create a Cross Entropy loss and a gradient.
End of explanation
predictor = nt.predict.gradient_descent(loss,
    g_dd, train['label'], learning_rate=learning_rate, momentum=mass)
fx_train = f(params, train['image'])
Explanation: Create a momentum predictor and initialize it.
End of explanation
print ('Time\tLoss\tLinear Loss')
X, Y = train['image'], train['label']
predictions = predictor(training_steps, fx_train)
for i in training_steps:
  params = get_params(state)
  state = opt_apply(i, grad_loss(params, X, Y), state)
  if i % print_every == 0:
    exact_loss = loss(f(params, X), Y)
    linear_loss = loss(predictions[i], Y)
    print('{:.0f}\t{:.4f}\t{:.4f}'.format(i, exact_loss, linear_loss))
Explanation: Train the network.
End of explanation
<END_TASK>