code: def main(unused_argv):
      # Build data pipelines.
      print("Loading data.")
      x_train, y_train, x_test, y_test = datasets.get_dataset(
          "cifar10", _TRAIN_SIZE, _TEST_SIZE
      )
  
      # Build the infinite network.
      _, _, kernel_fn = stax.serial(
          stax.Dense(1, 2.0, 0.05),
          stax.Relu(),
          stax.Dense(1, 2.0, 0.05),
      )
  
      # Optionally, compute the kernel in batches, in parallel.
      kernel_fn = nt.batch(
          kernel_fn, device_count=0, batch_size=_BATCH_SIZE
      )
  
      start = time.time()
      # Bayesian and infinite-time gradient descent inference with infinite network.
      predict_fn = nt.predict.gradient_descent_mse_ensemble(
          kernel_fn, x_train, y_train, diag_reg=1e-3
      )
      fx_test_nngp, fx_test_ntk = predict_fn(x_test=x_test)
      fx_test_nngp.block_until_ready()
      fx_test_ntk.block_until_ready()
  
      duration = time.time() - start
      print(
          "Kernel construction and inference done in %s seconds."
          % duration
      )
  
      # Print out accuracy and loss for infinite network predictions.
      loss = lambda fx, y_hat: 0.5 * jnp.mean((fx - y_hat) ** 2)
      util.print_summary("NNGP test", y_test, fx_test_nngp, None, loss)
      util.print_summary("NTK test", y_test, fx_test_ntk, None, loss)
  
file_path: neural-tangents/examples/infinite_fcn.py
imports: import time
  from absl import app
  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