code: def main(*args, use_dummy_data: bool = False, **kwargs) -> None:
      # Mask all padding with this value.
      mask_constant = 100.0
  
      if use_dummy_data:
          x_train, y_train, x_test, y_test = _get_dummy_data(
              mask_constant
          )
      else:
          # Build data pipelines.
          print("Loading IMDb data.")
          x_train, y_train, x_test, y_test = datasets.get_dataset(
              name="imdb_reviews",
              n_train=_TRAIN_SIZE,
              n_test=_TEST_SIZE,
              do_flatten_and_normalize=False,
              data_dir=_IMDB_PATH,
              input_key="text",
          )
  
          # Embed words and pad / truncate sentences to a fixed size.
          x_train, x_test = datasets.embed_glove(
              xs=[x_train, x_test],
              glove_path=_GLOVE_PATH,
              max_sentence_length=_MAX_SENTENCE_LENGTH,
              mask_constant=mask_constant,
          )
  
      # Build the infinite network.
      # Not using the finite model, hence width is set to 1 everywhere.
      _, _, kernel_fn = stax.serial(
          stax.Conv(
              out_chan=1,
              filter_shape=(9,),
              strides=(1,),
              padding="VALID",
          ),
          stax.Relu(),
          stax.GlobalSelfAttention(
              n_chan_out=1,
              n_chan_key=1,
              n_chan_val=1,
              pos_emb_type="SUM",
              W_pos_emb_std=1.0,
              pos_emb_decay_fn=lambda d: 1 / (1 + d**2),
              n_heads=1,
          ),
          stax.Relu(),
          stax.GlobalAvgPool(),
          stax.Dense(out_dim=1),
      )
  
      # Optionally, compute the kernel in batches, in parallel.
      kernel_fn = nt.batch(
          kernel_fn, device_count=-1, batch_size=_BATCH_SIZE
      )
  
      start = time.time()
      # Bayesian and infinite-time gradient descent inference with infinite network.
      predict = nt.predict.gradient_descent_mse_ensemble(
          kernel_fn=kernel_fn,
          x_train=x_train,
          y_train=y_train,
          diag_reg=1e-6,
          mask_constant=mask_constant,
      )
  
      fx_test_nngp, fx_test_ntk = predict(
          x_test=x_test, get=("nngp", "ntk")
      )
  
      fx_test_nngp.block_until_ready()
      fx_test_ntk.block_until_ready()
  
      duration = time.time() - start
      print(
          f"Kernel construction and inference done in {duration} seconds."
      )
  
      # 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/imdb.py
imports: import time
  from absl import app
  from jax import random
  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() without return types