code: def testPredictND(self):
      n_chan = 6
      key = random.PRNGKey(1)
      im_shape = (5, 4, 3)
      n_train = 2
      n_test = 2
      x_train = random.normal(key, (n_train,) + im_shape)
      y_train = random.uniform(key, (n_train, 3, 2, n_chan))
      init_fn, apply_fn, _ = stax.Conv(n_chan, (3, 2), (1, 2))
      _, params = init_fn(key, x_train.shape)
      fx_train_0 = apply_fn(params, x_train)
  
      for trace_axes in [
          (),
          (-1,),
          (-2,),
          (-3,),
          (0, 1),
          (2, 3),
          (2,),
          (1, 3),
          (0, -1),
          (0, 0, -3),
          (0, 1, 2, 3),
          (0, 1, -1, 2),
      ]:
          for ts in [None, jnp.arange(6).reshape((2, 3))]:
              for x in [None, "x_test"]:
                  with self.subTest(trace_axes=trace_axes, ts=ts, x=x):
                      t_shape = ts.shape if ts is not None else ()
                      y_test_shape = (
                          t_shape + (n_test,) + y_train.shape[1:]
                      )
                      y_train_shape = t_shape + y_train.shape
                      x = (
                          x
                          if x is None
                          else random.normal(key, (n_test,) + im_shape)
                      )
                      fx_test_0 = (
                          None if x is None else apply_fn(params, x)
                      )
  
                      kernel_fn = nt.empirical_kernel_fn(
                          apply_fn, trace_axes=trace_axes
                      )
  
                      kernel_fn = jit(kernel_fn, static_argnames="get")
                      ntk_train_train = kernel_fn(
                          x_train, None, "ntk", params
                      )
                      if x is not None:
                          ntk_test_train = kernel_fn(
                              x, x_train, "ntk", params
                          )
  
                      loss = lambda x, y: 0.5 * jnp.mean(x - y) ** 2
                      predict_fn_mse = predict.gradient_descent_mse(
                          ntk_train_train,
                          y_train,
                          trace_axes=trace_axes,
                      )
  
                      predict_fn_mse_ensemble = (
                          predict.gradient_descent_mse_ensemble(
                              kernel_fn,
                              x_train,
                              y_train,
                              trace_axes=trace_axes,
                              params=params,
                          )
                      )
  
                      if x is None:
                          p_train_mse = predict_fn_mse(ts, fx_train_0)
                      else:
                          p_train_mse, p_test_mse = predict_fn_mse(
                              ts, fx_train_0, fx_test_0, ntk_test_train
                          )
                          self.assertAllClose(
                              y_test_shape, p_test_mse.shape
                          )
                      self.assertAllClose(
                          y_train_shape, p_train_mse.shape
                      )  # pytype: disable=attribute-error  # jax-ndarray
  
                      (
                          p_nngp_mse_ens,
                          p_ntk_mse_ens,
                      ) = predict_fn_mse_ensemble(
                          ts, x, ("nngp", "ntk"), compute_cov=True
                      )
                      ref_shape = (
                          y_train_shape if x is None else y_test_shape
                      )
                      self.assertAllClose(
                          ref_shape, p_ntk_mse_ens.mean.shape
                      )
                      self.assertAllClose(
                          ref_shape, p_nngp_mse_ens.mean.shape
                      )
  
                      if ts is not None:
                          predict_fn = predict.gradient_descent(
                              loss,
                              ntk_train_train,
                              y_train,
                              trace_axes=trace_axes,
                          )
  
                          if x is None:
                              p_train = predict_fn(ts, fx_train_0)
                          else:
                              p_train, p_test = predict_fn(
                                  ts,
                                  fx_train_0,
                                  fx_test_0,
                                  ntk_test_train,
                              )
                              self.assertAllClose(
                                  y_test_shape, p_test.shape
                              )
                          self.assertAllClose(
                              y_train_shape, p_train.shape
                          )
  
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 18 for calling the PredictTest obj's underlying member method code with example usage: obj.testPredictND() without return types