code: def gradient_descent_mse(
      k_train_train: jnp.ndarray,
      y_train: jnp.ndarray,
      learning_rate: float = 1.0,
      diag_reg: float = 0.0,
      diag_reg_absolute_scale: bool = False,
      trace_axes: Axes = (-1,),
  ) -> PredictFn:
      r"""Predicts the outcome of function space gradient descent training on MSE.
  
      Solves in closed form for the continuous-time version of gradient descent.
  
      Uses the closed-form solution for gradient descent on an MSE loss in function
      space detailed in [*,**] given a Neural Tangent or Neural Network Gaussian
      Process Kernel over the dataset. Given NNGP or NTK, this function will return
      a function that predicts the time evolution for function space points at
      arbitrary time[s] (training step[s]) `t`. Note that these time[s] (step[s])
      are continuous and are interpreted in units of the `learning_rate` so
      `absolute_time = learning_rate * t`, and the scales of `learning_rate` and `t`
      are interchangeable.
  
      Note that first invocation of the returned `predict_fn` will be slow and
      allocate a lot of memory for its whole lifetime, as either eigendecomposition
      (`t` is a scalar or an array) or Cholesky factorization (`t=None`) of
      `k_train_train` is performed and cached for future invocations (or both, if
      the function is called on both finite and infinite (`t=None`) times).
  
      [*] "`Neural Tangent Kernel: Convergence and Generalization in Neural Networks
      <https://arxiv.org/abs/1806.07572>`_"
  
      [**] "`Wide Neural Networks of Any Depth Evolve as Linear
      Models Under Gradient Descent <https://arxiv.org/abs/1902.06720>`_"
  
      Example:
        >>> import neural_tangents as nt
        >>> #
        >>> t = 1e-7
        >>> kernel_fn = nt.empirical_ntk_fn(f)
        >>> k_train_train = kernel_fn(x_train, None, params)
        >>> k_test_train = kernel_fn(x_test, x_train, params)
        >>> #
        >>> predict_fn = nt.predict.gradient_descent_mse(k_train_train, y_train)
        >>> #
        >>> fx_train_0 = f(params, x_train)
        >>> fx_test_0 = f(params, x_test)
        >>> #
        >>> fx_train_t, fx_test_t = predict_fn(t, fx_train_0, fx_test_0,
        >>>                                    k_test_train)
  
      Args:
        k_train_train:
          kernel on the training data. Must have the shape of
          `zip(y_train.shape, y_train.shape)` with `trace_axes` absent.
  
        y_train:
          targets for the training data.
  
        learning_rate:
          learning rate, step size.
  
        diag_reg:
          a scalar representing the strength of the diagonal regularization for
          `k_train_train`, i.e. computing `k_train_train + diag_reg * I` during
          Cholesky factorization or eigendecomposition.
  
        diag_reg_absolute_scale:
          `True` for `diag_reg` to represent regularization in absolute units,
          `False` to be `diag_reg * jnp.mean(jnp.trace(k_train_train))`.
  
        trace_axes:
          `f(x_train)` axes such that `k_train_train` lacks these pairs of
          dimensions and is to be interpreted as :math:`\Theta \otimes I`, i.e.
          block-diagonal along `trace_axes`. These can can be specified either to
          save space and compute, or to even improve approximation accuracy of the
          infinite-width or infinite-samples limit, since in these limits the
          covariance along channel / feature / logit axes indeed converges to a
          constant-diagonal matrix. However, if you target linearized dynamics of a
          specific finite-width network, `trace_axes=()` will yield most accurate
          result.
  
      Returns:
        A function of signature
        `predict_fn(t, fx_train_0, fx_test_0, k_test_train)` that
        returns output train [and test] set[s] predictions at time[s] `t`.
      """
      _, odd, first, _ = _get_axes(k_train_train)
      trace_axes = utils.canonicalize_axis(trace_axes, y_train)
      trace_axes = tuple(-y_train.ndim + a for a in trace_axes)
      n_t_axes, n_non_t_axes = len(trace_axes), y_train.ndim - len(
          trace_axes
      )
      last_t_axes = tuple(range(-n_t_axes, 0))
      non_t_axes = tuple(range(-y_train.ndim, -n_t_axes))
  
      @lru_cache(1)
      def get_predict_fn_inf():
          with jax.core.eval_context():
              solve = _get_cho_solve(
                  k_train_train, diag_reg, diag_reg_absolute_scale
              )
  
          def predict_fn_inf(fx_train_0, fx_test_0, k_test_train):
              fx_train_t = y_train.astype(k_train_train.dtype)
              if fx_test_0 is None:
                  return fx_train_t
  
              rhs = (
                  y_train
                  if fx_train_0 is None
                  else y_train - fx_train_0
              )
              dfx_test = jnp.tensordot(
                  k_test_train, solve(rhs, trace_axes), (odd, first)
              )
              dfx_test = jnp.moveaxis(dfx_test, last_t_axes, trace_axes)
              fx_test_t = fx_test_0 + dfx_test
  
              if fx_train_0 is None:
                  return fx_test_t
              return fx_train_t, fx_test_t
  
          return predict_fn_inf
  
      @lru_cache(1)
      def get_predict_fn_finite():
          with jax.core.eval_context():
              expm1_fn, inv_expm1_fn = _get_fns_in_eigenbasis(
                  k_train_train,
                  diag_reg,
                  diag_reg_absolute_scale,
                  (
                      _make_expm1_fn(y_train.size),
                      _make_inv_expm1_fn(y_train.size),
                  ),
              )
  
          rhs_shape = tuple(y_train.shape[a] for a in trace_axes)
  
          def predict_fn_finite(t, fx_train_0, fx_test_0, k_test_train):
              t = jnp.array(t) * learning_rate
              t_shape, t_ndim = t.shape, t.ndim
              first_t_axes = tuple(range(t_ndim))
              t = t.reshape((-1, 1))
  
              rhs = (
                  -y_train
                  if fx_train_0 is None
                  else fx_train_0 - y_train
              )
              rhs = jnp.moveaxis(rhs, trace_axes, last_t_axes).reshape(
                  (-1,) + rhs_shape
              )
              shape = t_shape + k_train_train.shape[1::2] + rhs_shape
  
              if fx_train_0 is not None:
                  dfx_train = expm1_fn(rhs, t).reshape(shape)
                  dfx_train = jnp.moveaxis(
                      dfx_train, last_t_axes, trace_axes
                  )
                  fx_train_t = (
                      jnp.expand_dims(fx_train_0, first_t_axes)
                      + dfx_train
                  )
  
              if fx_test_0 is not None:
                  dfx_test = inv_expm1_fn(rhs, t).reshape(shape)
                  dfx_test = jnp.tensordot(
                      k_test_train, dfx_test, (odd, non_t_axes)
                  )
                  dfx_test = jnp.moveaxis(
                      dfx_test,
                      tuple(range(n_non_t_axes, n_non_t_axes + t_ndim))
                      + last_t_axes,
                      tuple(range(t_ndim)) + trace_axes,
                  )
                  fx_test_t = (
                      jnp.expand_dims(fx_test_0, first_t_axes)
                      + dfx_test
                  )
  
              if fx_train_0 is not None and fx_test_0 is not None:
                  return fx_train_t, fx_test_t
              if fx_test_0 is None:
                  return fx_train_t
              return fx_test_t
  
          return predict_fn_finite
  
      def predict_fn(
          t: Optional[ArrayOrScalar] = None,
          fx_train_0: ArrayOrScalar = 0.0,
          fx_test_0: Optional[ArrayOrScalar] = None,
          k_test_train: Optional[jnp.ndarray] = None,
      ) -> Union[jnp.ndarray, tuple[jnp.ndarray, jnp.ndarray]]:
          """Return output predictions on train [and test] set[s] at time[s] `t`.
  
          Args:
            t:
              a scalar of array of scalars of any shape. `t=None` is treated as
              infinity and returns the same result as `t=jnp.inf`, but is computed
              using identity or linear solve for train and test predictions
              respectively instead of eigendecomposition, saving time and precision.
              Equivalent of training steps (but can be fractional).
  
            fx_train_0:
              output of the network at `t == 0` on the training set. `fx_train_0=None`
              means to not compute predictions on the training set.
  
            fx_test_0:
              output of the network at `t == 0` on the test set. `fx_test_0=None`
              means to not compute predictions on the test set.
  
            k_test_train:
              kernel relating test data with training data. Must have the shape of
              `zip(y_test.shape, y_train.shape)` with `trace_axes` absent. Pass
              `k_test_train=None` if you only need non-regularized (`diag_reg=0`)
              predictions on the training set. For regularized train-set predictions,
              pass `k_test_train=k_train_train`.
  
          Returns:
            `fx_train_t` or `(fx_train_t, fx_test_t)` if `fx_test_0 != None` with
            potentially additional leading time dimensions matching `t.shape`.
  
          Raises:
            ValueError: if `fx_test_0` is not `None`, but `k_test_train` is `None`.
          """
          _check_inputs(fx_train_0, fx_test_0, k_test_train)
  
          # Infinite time
          if t is None:
              return get_predict_fn_inf()(
                  fx_train_0, fx_test_0, k_test_train
              )
  
          # Finite time
          return get_predict_fn_finite()(
              t, fx_train_0, fx_test_0, k_test_train
          )
  
      return predict_fn
  
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 1 for calling a global function with example usage: gradient_descent_mse(k_train_train, y_train, learning_rate, diag_reg, diag_reg_absolute_scale, trace_axes) and returns: predict_fn, predict_fn_inf, predict_fn_finite, get_predict_fn_finite, fx_train_t,  fx_test_t, fx_test_t, get_predict_fn_inf, fx_train_t, fx_test_t, fx_train_t,  fx_test_t, fx_train_t