code: def main(unused_argv):
      # Build data pipelines.
      print("Loading data.")
      x_train, y_train, x_test, y_test = datasets.get_dataset(
          "mnist", _TRAIN_SIZE, _TEST_SIZE
      )
  
      # Build the network
      init_fn, apply_fn, _ = stax.serial(
          stax.Dense(512, 1.0, 0.05),
          stax.Erf(),
          stax.Dense(10, 1.0, 0.05),
      )
  
      key = random.PRNGKey(0)
      _, params = init_fn(key, (-1, 784))
  
      # Create and initialize an optimizer.
      opt_init, opt_apply, get_params = optimizers.sgd(_LEARNING_RATE)
      state = opt_init(params)
  
      # Create an mse loss function and a gradient function.
      loss = lambda fx, y_hat: 0.5 * jnp.mean((fx - y_hat) ** 2)
      grad_loss = jit(
          grad(lambda params, x, y: loss(apply_fn(params, x), y))
      )
  
      # Create an MSE predictor to solve the NTK equation in function space.
      ntk = nt.batch(
          nt.empirical_ntk_fn(apply_fn, vmap_axes=0),
          batch_size=64,
          device_count=0,
      )
      g_dd = ntk(x_train, None, params)
      g_td = ntk(x_test, x_train, params)
      predictor = nt.predict.gradient_descent_mse(g_dd, y_train)
  
      # Get initial values of the network in function space.
      fx_train = apply_fn(params, x_train)
      fx_test = apply_fn(params, x_test)
  
      # Train the network.
      train_steps = int(_TRAIN_TIME // _LEARNING_RATE)
      print("Training for {} steps".format(train_steps))
  
      for i in range(train_steps):
          params = get_params(state)
          state = opt_apply(
              i, grad_loss(params, x_train, y_train), state
          )
  
      # Get predictions from analytic computation.
      print("Computing analytic prediction.")
      fx_train, fx_test = predictor(
          _TRAIN_TIME, fx_train, fx_test, g_td
      )
  
      # Print out summary data comparing the linear / nonlinear model.
      util.print_summary(
          "train", y_train, apply_fn(params, x_train), fx_train, loss
      )
      util.print_summary(
          "test", y_test, apply_fn(params, x_test), fx_test, loss
      )
  
file_path: neural-tangents/examples/function_space.py
imports: from absl import app
  from jax import grad
  from jax import jit
  from jax import random
  from jax.example_libraries import optimizers
  import jax.numpy
  import neural_tangents
  from neural_tangents import stax
  from examples import datasets
  from examples import util
label_desc: Use image node_id 1 for calling a global function with example usage: main(unused_argv) without return types