code: def testTrainedEnsemblePredCov(
      self, train_size, test_size, input_shape, out_logits
  ):
      training_steps = 1000
      learning_rate = 0.1
      ensemble_size = 1024
  
      init_fn, apply_fn, kernel_fn = stax.serial(
          stax.Dense(128, W_std=1.2, b_std=0.05),
          stax.Erf(),
          stax.Dense(out_logits, W_std=1.2, b_std=0.05),
      )
  
      opt_init, opt_update, get_params = optimizers.sgd(learning_rate)
      opt_update = jit(opt_update)
  
      train_shape = (train_size, *input_shape)
      test_shape = (test_size, *input_shape)
  
      key, x_test, x_train, y_train = self._get_inputs(
          out_logits, test_shape, train_shape
      )
      predict_fn_mse_ens = predict.gradient_descent_mse_ensemble(
          kernel_fn,
          x_train,
          y_train,
          learning_rate=learning_rate,
          diag_reg=0.0,
      )
  
      train = (x_train, y_train)
      ensemble_key = random.split(key, ensemble_size)
  
      loss = jit(
          lambda params, x, y: 0.5
          * jnp.mean((apply_fn(params, x) - y) ** 2)
      )
      grad_loss = jit(
          lambda state, x, y: grad(loss)(get_params(state), x, y)
      )
  
      def train_network(key):
          _, params = init_fn(key, (-1,) + train_shape[1:])
          opt_state = opt_init(params)
          for i in range(training_steps):
              opt_state = opt_update(
                  i, grad_loss(opt_state, *train), opt_state
              )
  
          return get_params(opt_state)
  
      params = vmap(train_network)(ensemble_key)
      tol = 0.08
  
      for x in [None, "x_test"]:
          with self.subTest(x=x):
              x = x if x is None else x_test
              x_fin = x_train if x is None else x_test
              ensemble_fx = vmap(apply_fn, (0, None))(params, x_fin)
  
              mean_emp = jnp.mean(ensemble_fx, axis=0, keepdims=True)
              mean_subtracted = ensemble_fx - mean_emp
              cov_emp = jnp.einsum(  # pytype: disable=wrong-arg-types  # jnp-type
                  "ijk,ilk->jl",
                  mean_subtracted,
                  mean_subtracted,
                  optimize=True,
              ) / (
                  mean_subtracted.shape[0] * mean_subtracted.shape[-1]
              )
  
              ntk = predict_fn_mse_ens(
                  training_steps, x, "ntk", compute_cov=True
              )
              self.assertAllClose(
                  ravel_pytree(mean_emp)[0],
                  ravel_pytree(ntk.mean)[0],
                  rtol=tol,
                  atol=tol,
              )
              self.assertAllClose(
                  cov_emp, ntk.covariance, rtol=tol, atol=tol
              )
  
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 14 for calling the PredictTest obj's underlying member method code with example usage: obj.testTrainedEnsemblePredCov(train_size, test_size, input_shape, out_logits) and returns: get_params