code: def _test_multi_step(
      self, predictor, fx_train_0, fx_test_0, g_td, momentum
  ):
      # Test multi-time prediction
      ts = jnp.arange(6).reshape((2, 1, 3))
  
      fx_train_single, fx_test_single = predictor(
          ts, fx_train_0, fx_test_0, g_td
      )
  
      fx_train_concat, fx_test_concat = [], []
      for t in ts.ravel():
          fx_train_concat_t, fx_test_concat_t = predictor(
              t, fx_train_0, fx_test_0, g_td
          )
          fx_train_concat += [fx_train_concat_t]
          fx_test_concat += [fx_test_concat_t]
      fx_train_concat = jnp.stack(fx_train_concat).reshape(
          ts.shape + fx_train_single.shape[ts.ndim :]
      )
      fx_test_concat = jnp.stack(fx_test_concat).reshape(
          ts.shape + fx_test_single.shape[ts.ndim :]
      )
  
      self.assertAllClose(fx_train_concat, fx_train_single)
      self.assertAllClose(fx_test_concat, fx_test_single)
  
      if momentum is not None:
          state_0 = predict.ODEState(
              fx_train_0, fx_test_0
          )  # pytype:disable=wrong-arg-count
          t_1 = (0, 0, 2)
          state_1 = predictor(ts[t_1], state_0, None, g_td)
          self.assertAllClose(fx_train_single[t_1], state_1.fx_train)
          self.assertAllClose(fx_test_single[t_1], state_1.fx_test)
  
          t_max = (-1,) * ts.ndim
          state_max = predictor(
              ts[t_max] - ts[t_1], state_1, None, g_td
          )
          self.assertAllClose(
              fx_train_single[t_max], state_max.fx_train
          )
          self.assertAllClose(fx_test_single[t_max], state_max.fx_test)
  
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 3 for calling the PredictTest obj's underlying member method code with example usage: obj._test_multi_step(predictor, fx_train_0, fx_test_0, g_td, momentum) without return types