code: def _test_zero_time(
      self, predictor, fx_train_0, fx_test_0, g_td, momentum
  ):
      fx_train_t0, fx_test_t0 = predictor(
          0.0, fx_train_0, fx_test_0, g_td
      )
      self.assertAllClose(fx_train_0, fx_train_t0)
      self.assertAllClose(fx_test_0, fx_test_t0)
      fx_train_only_t0 = predictor(0.0, fx_train_0, None, g_td)
      self.assertAllClose(fx_train_0, fx_train_only_t0)
  
      if momentum is not None:
          # Test state-based prediction
          state_0 = predict.ODEState(
              fx_train_0, fx_test_0
          )  # pytype:disable=wrong-arg-count
          state_t0 = predictor(0.0, state_0, None, g_td)
          self.assertAllClose(state_0.fx_train, state_t0.fx_train)
          self.assertAllClose(state_0.fx_test, state_t0.fx_test)
  
          state_train_only_0 = predict.ODEState(
              fx_train_0
          )  # pytype:disable=wrong-arg-count
          state_train_only_t0 = predictor(0.0, state_0, None, g_td)
          self.assertAllClose(
              state_train_only_0.fx_train, state_train_only_t0.fx_train
          )
  
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 PredictTest obj's underlying member method code with example usage: obj._test_zero_time(predictor, fx_train_0, fx_test_0, g_td, momentum) without return types