code: def testMaxLearningRate(
      self,
      train_size,
      input_shape,
      network,
      out_logits,
      kernel_type,
      lr_factor,
      momentum,
  ):
      key = random.PRNGKey(0)
  
      key, split = random.split(key)
      if len(input_shape) == 1:
          train_shape = (train_size * 5, input_shape[0] * 10)
      else:
          train_shape = (16, 8, 8, 3)
      x_train = random.normal(split, train_shape)
  
      key, split = random.split(key)
      y_train = jnp.array(
          random.bernoulli(split, shape=(train_shape[0], out_logits)),
          jnp.float32,
      )
  
      # Regress to an MSE loss.
      loss = lambda params, x: 0.5 * jnp.mean(
          (f(params, x) - y_train) ** 2
      )
      grad_loss = jit(grad(loss))
  
      def get_loss(opt_state):
          return loss(get_params(opt_state), x_train)
  
      steps = 30
  
      fn_and_kernel = KERNELS[kernel_type]
      params, f, ntk = fn_and_kernel(
          key, train_shape[1:], network, out_logits
      )
      g_dd = ntk(x_train, None, "ntk")
  
      step_size = (
          predict.max_learning_rate(
              g_dd, y_train_size=y_train.size, momentum=momentum
          )
          * lr_factor
      )
      opt_init, opt_update, get_params = optimizers.momentum(
          step_size, mass=momentum
      )
  
      opt_state = opt_init(params)
  
      init_loss = get_loss(opt_state)
  
      for i in range(steps):
          params = get_params(opt_state)
          opt_state = opt_update(
              i, grad_loss(params, x_train), opt_state
          )
  
      trained_loss = get_loss(opt_state)
      loss_ratio = trained_loss / (init_loss + 1e-12)
      if lr_factor < 1.0:
          self.assertLess(loss_ratio, 0.1)
      elif lr_factor == 1:
          # At the threshold, the loss decays slowly
          self.assertLess(loss_ratio, 1.0)
      if lr_factor > 2.0:
          if not math.isnan(loss_ratio):
              self.assertGreater(loss_ratio, 10.0)
  
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 19 for calling the PredictTest obj's underlying member method code with example usage: obj.testMaxLearningRate(train_size, input_shape, network, out_logits, kernel_type, lr_factor, momentum) and returns: loss