code: def test_kwargs(self, do_batch, mode):
      rng = random.PRNGKey(1)
  
      x_train = random.normal(rng, (8, 7, 10))
      x_test = random.normal(rng, (4, 7, 10))
      y_train = random.normal(rng, (8, 1))
  
      rng_train, rng_test = random.split(rng, 2)
  
      pattern_train = random.normal(rng, (8, 7, 7))
      pattern_test = random.normal(rng, (4, 7, 7))
  
      diag_reg = 1e-4
  
      if jax.default_backend() == "tpu":
          atol = 3e-3
          rtol = 4e-2
          width = 256
      else:
          atol = 5e-4
          rtol = 1e-2
          width = 64
  
      init_fn, apply_fn, kernel_fn = stax.serial(
          stax.Dense(width, W_std=2**0.5),
          stax.Relu(),
          stax.Dropout(rate=0.7),
          stax.Aggregate(),
          stax.GlobalAvgPool(),
          stax.Dense(width),
      )
  
      kw_dd = dict(pattern=(pattern_train, pattern_train))
      kw_td = dict(pattern=(pattern_test, pattern_train))
      kw_tt = dict(pattern=(pattern_test, pattern_test))
  
      if mode == "mc":
          kernel_fn = nt.monte_carlo_kernel_fn(
              init_fn, apply_fn, rng, 2, batch_size=2 if do_batch else 0
          )
  
      elif mode == "empirical":
          kernel_fn = nt.empirical_kernel_fn(apply_fn)
          if do_batch:
              raise absltest.SkipTest(
                  "Batching of empirical kernel is not "
                  "implemented with keyword arguments."
              )
  
          for kw in (kw_dd, kw_td, kw_tt):
              kw.update(
                  dict(
                      params=init_fn(rng, x_train.shape)[1],
                      get=("nngp", "ntk"),
                  )
              )
  
          kw_dd.update(dict(rng=(rng_train, None)))
          kw_td.update(dict(rng=(rng_test, rng_train)))
          kw_tt.update(dict(rng=(rng_test, None)))
  
      elif mode == "analytic":
          if do_batch:
              kernel_fn = nt.batch(kernel_fn, batch_size=2)
  
      else:
          raise ValueError(mode)
  
      k_dd = kernel_fn(x_train, None, **kw_dd)
      k_td = kernel_fn(x_test, x_train, **kw_td)
      k_tt = kernel_fn(x_test, None, **kw_tt)
  
      # Infinite time NNGP/NTK.
      predict_fn_gp = predict.gp_inference(
          k_dd, y_train, diag_reg=diag_reg
      )
      out_gp = predict_fn_gp(k_test_train=k_td, k_test_test=k_tt.nngp)
  
      if mode == "empirical":
          for kw in (kw_dd, kw_td, kw_tt):
              kw.pop("get")
  
      predict_fn_ensemble = predict.gradient_descent_mse_ensemble(
          kernel_fn, x_train, y_train, diag_reg=diag_reg, **kw_dd
      )
      out_ensemble = predict_fn_ensemble(
          x_test=x_test, compute_cov=True, **kw_tt
      )
      self.assertAllClose(out_gp, out_ensemble)
  
      # Finite time NTK test.
      predict_fn_mse = predict.gradient_descent_mse(k_dd.ntk, y_train)
      out_mse = predict_fn_mse(
          t=1.0, fx_train_0=None, fx_test_0=0.0, k_test_train=k_td.ntk
      )
      out_ensemble = predict_fn_ensemble(
          t=1.0, get="ntk", x_test=x_test, compute_cov=False, **kw_tt
      )
      self.assertAllClose(out_mse, out_ensemble, atol=atol, rtol=rtol)
  
      # Finite time NTK train.
      out_mse = predict_fn_mse(
          t=0.5, fx_train_0=0.0, fx_test_0=None, k_test_train=k_td.ntk
      )
      out_ensemble = predict_fn_ensemble(
          t=0.5, get="ntk", x_test=None, compute_cov=False, **kw_dd
      )
      self.assertAllClose(out_mse, out_ensemble, atol=atol, rtol=rtol)
  
      # Finite time NNGP test.
      predict_fn_mse = predict.gradient_descent_mse(k_dd.nngp, y_train)
      out_mse = predict_fn_mse(
          t=1.0, fx_train_0=None, fx_test_0=0.0, k_test_train=k_td.nngp
      )
      out_ensemble = predict_fn_ensemble(
          t=1.0, get="nngp", x_test=x_test, compute_cov=False, **kw_tt
      )
      self.assertAllClose(out_mse, out_ensemble, atol=atol, rtol=rtol)
  
      # Finite time NNGP train.
      out_mse = predict_fn_mse(
          t=5.0, fx_train_0=0.0, fx_test_0=None, k_test_train=k_td.nngp
      )
      out_ensemble = predict_fn_ensemble(
          t=5.0, get="nngp", x_test=None, compute_cov=False, **kw_dd
      )
      self.assertAllClose(out_mse, out_ensemble, atol=atol, rtol=rtol)
  
imports: import math
  from absl.testing import absltest
  from jax import grad
  from jax import jit
  from jax import random
  from jax import vmap
  from jax import config
  from jax.example_libraries import optimizers
  from jax.flatten_util import ravel_pytree
  import jax.numpy
  import jax.tree_util
  import neural_tangents
  from neural_tangents import predict, stax
  from neural_tangents._src.predict import _is_on_cpu
  from tests import test_utils
label_desc: Use image node_id 1 for calling the PredictKwargsTest obj's underlying member method code with example usage: obj.test_kwargs(do_batch, mode) without return types