code: def max_learning_rate(
      ntk_train_train: jnp.ndarray,
      y_train_size: Optional[int] = None,
      momentum=0.0,
      eps: float = 1e-12,
  ) -> float:
      r"""Computes the maximal feasible learning rate for infinite width NNs.
  
      The network is assumed to be trained using mini-/full-batch GD + momentum
      with mean squared loss. The loss is assumed to have the form
      `1/(2 * batch_size * output_size) \|f(train_x) - train_y\|^2`. For vanilla SGD
      (i.e. `momentum = 0`) the maximal feasible learning rate is the largest `\eta`
      such that the operator `(I - \eta / (batch_size * output_size) * NTK)` is a
      contraction, which is `2 * batch_size * output_size * lambda_max(NTK)`. When
      `momentum > 0`, we use
      `2 * (1 + momentum) * batch_size * output_size * lambda_max(NTK)` (see
      *The Dynamics of Momentum* section in
      "`Why Momentum Really Works <https://distill.pub/2017/momentum/>`_").
  
      Args:
        ntk_train_train:
          analytic or empirical NTK on the training data.
  
        y_train_size:
          total training set output size, i.e.
          `f(x_train).size ==  y_train.size`. If `output_size=None` it is inferred
          from `ntk_train_train.shape` assuming `trace_axes=()`.
  
        momentum:
          The `momentum` for momentum optimizers.
  
        eps:
          a float to avoid zero divisor.
  
      Returns:
        The maximal feasible learning rate for infinite width NNs.
      """
      ntk_train_train = utils.make_2d(ntk_train_train)
      factor = (
          ntk_train_train.shape[0]
          if y_train_size is None
          else y_train_size
      )  # pytype: disable=attribute-error  # jax-ndarray
  
      if _is_on_cpu(ntk_train_train):
          max_eva = sp.linalg.eigvalsh(
              ntk_train_train,
              eigvals=(
                  ntk_train_train.shape[0]
                  - 1,  # pytype: disable=attribute-error  # jax-ndarray
                  ntk_train_train.shape[0] - 1,
              ),
          )[
              -1
          ]  # pytype: disable=attribute-error  # jax-ndarray
      else:
          max_eva = jnp.linalg.eigvalsh(ntk_train_train)[-1]
      lr = 2 * (1 + momentum) * factor / (max_eva + eps)
      return lr
  
imports: import collections
  from functools import lru_cache
  from typing import Any, Callable, Generator, Iterable, NamedTuple, Optional, Protocol, Union
  import jax
  from jax import grad
  from jax.experimental import ode
  import jax.numpy
  import jax.scipy
  from jax.tree_util import tree_all
  from jax.tree_util import tree_map
  import numpy
  import scipy
  from .utils import dataclasses
  from .utils import utils
  from .utils.typing import Axes
  from .utils.typing import Get
  from .utils.typing import KernelFn
label_desc: Use image node_id 5 for calling a global function with example usage: max_learning_rate(ntk_train_train, y_train_size, momentum, eps) and returns: lr