docstr: Test that the NTK and NNGP agree at t=0.
docstr_tok: ["Test","that","the","NTK","and","NNGP","agree","at","t=0","."]
code: def testZeroTimeAgreement(
      self,
      train_size,
      test_size,
      input_shape,
      network,
      out_logits,
  ):
      """Test that the NTK and NNGP agree at t=0."""
      train_shape = (train_size, *input_shape)
      test_shape = (test_size, *input_shape)
      _, x_test, x_train, y_train = self._get_inputs(
          out_logits, test_shape, train_shape
      )
      _, _, ker_fun = _build_network(
          train_shape[1:], network, out_logits
      )
  
      reg = 1e-7
      predictor = predict.gradient_descent_mse_ensemble(
          ker_fun, x_train, y_train, diag_reg=reg
      )
  
      for x in (None, "x_test"):
          with self.subTest(x=x):
              x = x if x is None else x_test
              zero = predictor(
                  t=0.0, x_test=x, get=("NTK", "NNGP"), compute_cov=True
              )
              if x is None:
                  k = ker_fun(x_train, None, get="nngp")
                  ref = (jnp.zeros_like(y_train, k.dtype), k)
              else:
                  ref = (
                      jnp.zeros((test_shape[0], out_logits)),
                      ker_fun(x_test, None, get="nngp"),
                  )
  
              self.assertAllClose((ref,) * 2, zero, check_dtypes=False)
              if x is None:
                  zero_x = predictor(
                      t=0.0,
                      x_test=x_train,
                      get=("NTK", "NNGP"),
                      compute_cov=True,
                  )
                  self.assertAllClose((ref,) * 2, zero_x)
  
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 10 for calling the PredictTest obj's underlying member method code with example usage: obj.testZeroTimeAgreement(train_size, test_size, input_shape, network, out_logits) without return types