code: def gradient_descent(
      loss: Callable[[jnp.ndarray, jnp.ndarray], float],
      k_train_train: jnp.ndarray,
      y_train: jnp.ndarray,
      learning_rate: float = 1.0,
      momentum: Optional[float] = None,
      trace_axes: Axes = (-1,),
  ) -> PredictFnODE:
      r"""Predicts the outcome of function space training using gradient descent.
  
      Uses an ODE solver. If `momentum != None`, solves a continuous-time version of
      gradient descent with momentum.
  
      .. note::
        We use standard momentum as opposed to Nesterov momentum.
  
      Solves the function space ODE for [momentum] gradient descent with a given
      `loss` (detailed in "`Wide Neural Networks of Any Depth Evolve as Linear
      Models Under Gradient Descent <https://arxiv.org/abs/1902.06720>`_".) given a
      Neural Tangent Kernel[s] over the dataset[s] at arbitrary time[s] (step[s])
      `t`. Note that for gradient descent `absolute_time = learning_rate * t` and
      the scales of the learning rate and query step[s] `t` are interchangeable.
      However, the momentum gradient descent ODE is solved in the units of
      `learning_rate**0.5`, and therefore `absolute_time = learning_rate**0.5 * t`,
      hence the `learning_rate` and training time[s] (step[s]) `t` scales are not
      interchangeable.
  
      Example:
        >>> import neural_tangents as nt
        >>> #
        >>> t = 1e-7
        >>> learning_rate = 1e-2
        >>> momentum = 0.9
        >>> #
        >>> kernel_fn = nt.empirical_ntk_fn(f)
        >>> k_test_train = kernel_fn(x_test, x_train, params)
        >>> #
        >>> from jax.nn import log_softmax
        >>> cross_entropy = lambda fx, y_hat: -jnp.mean(log_softmax(fx) * y_hat)
        >>> predict_fn = nt.redict.gradient_descent(
        >>>     cross_entropy, k_train_train, y_train, learning_rate, momentum)
        >>> #
        >>> 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:
        loss:
          a loss function whose signature is `loss(f(x_train), y_train)`. Note:
          the loss function should treat the batch and output dimensions
          symmetrically.
  
        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.
  
        momentum:
          momentum scalar.
  
        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 that returns output train [and test] set[s] predictions at
        time[s] `t`.
      """
      _, odd, _, _ = _get_axes(k_train_train)
      trace_axes = utils.canonicalize_axis(trace_axes, y_train)
      non_t_axes = tuple(
          a for a in range(y_train.ndim) if a not in trace_axes
      )
      last_t_axes = range(-len(trace_axes), 0)
  
      dtype = k_train_train.dtype
      grad_loss = grad(lambda fx: loss(fx, y_train))
  
      if momentum is not None:
          learning_rate **= 0.5
          momentum = (momentum - 1.0) / learning_rate
  
      def get_state_0(fx_train_or_state_0, fx_test_0, fx_test_shape):
          if isinstance(fx_train_or_state_0, ODEState):
              fx_train_0 = fx_train_or_state_0.fx_train
              fx_test_0 = fx_train_or_state_0.fx_test
              qx_train_0 = fx_train_or_state_0.qx_train
              qx_test_0 = fx_train_or_state_0.qx_test
          else:
              fx_train_0 = fx_train_or_state_0
              qx_train_0 = qx_test_0 = None
  
          if fx_train_0 is None:
              fx_train_0 = jnp.zeros_like(y_train, dtype)
          else:
              fx_train_0 = jnp.broadcast_to(fx_train_0, y_train.shape)
  
          if fx_test_0 is not None:
              fx_test_0 = jnp.broadcast_to(fx_test_0, fx_test_shape)
  
          if momentum is None:
              if qx_train_0 is not None or qx_test_0 is not None:
                  raise ValueError(
                      "Got passed momentum state variables, while "
                      "`momentum is None`."
                  )
          else:
              qx_train_0 = (
                  jnp.zeros_like(y_train, dtype)
                  if qx_train_0 is None
                  else jnp.broadcast_to(qx_train_0, y_train.shape)
              )
              qx_test_0 = (
                  None
                  if fx_test_0 is None
                  else (
                      jnp.zeros(fx_test_shape, dtype)
                      if qx_test_0 is None
                      else jnp.broadcast_to(qx_test_0, fx_test_shape)
                  )
              )
  
          return ODEState(
              fx_train_0, fx_test_0, qx_train_0, qx_test_0
          )  # pytype: disable=wrong-arg-count
  
      def get_dstate_dt(k_test_train):
          def dstate_dt(state_t: ODEState, unused_t) -> ODEState:
              fx_train_t, fx_test_t, qx_train_t, qx_test_t = (
                  state_t.fx_train,
                  state_t.fx_test,
                  state_t.qx_train,
                  state_t.qx_test,
              )
  
              dy_df_t = grad_loss(fx_train_t)
  
              fx_train_t = -jnp.moveaxis(
                  jnp.tensordot(
                      k_train_train, dy_df_t, (odd, non_t_axes)
                  ),
                  last_t_axes,
                  trace_axes,
              )
              if fx_test_t is not None:
                  fx_test_t = -jnp.moveaxis(
                      jnp.tensordot(
                          k_test_train, dy_df_t, (odd, non_t_axes)
                      ),
                      last_t_axes,
                      trace_axes,
                  )
  
              if momentum is None:
                  return ODEState(fx_train_t, fx_test_t)
  
              fx_train_t += momentum * qx_train_t
              if qx_test_t is not None:
                  fx_test_t += momentum * qx_test_t
  
              return ODEState(
                  qx_train_t, qx_test_t, fx_train_t, fx_test_t
              )  # pytype: disable=wrong-arg-count
  
          return dstate_dt
  
      def predict_fn(
          t: Optional[ArrayOrScalar] = None,
          fx_train_or_state_0: Union[ArrayOrScalar, ODEState] = 0.0,
          fx_test_0: Optional[ArrayOrScalar] = None,
          k_test_train: Optional[jnp.ndarray] = None,
      ) -> Union[
          jnp.ndarray, tuple[jnp.ndarray, jnp.ndarray], ODEState
      ]:
          """Return output predictions on train [and test] set[s] at time[s] `t`.
  
          Args:
            t:
              a scalar or array of scalars of any shape in strictly increasing order.
              `t=None` is equivalent to `t=jnp.inf` and may not converge. Equivalent
              of training steps (but can be fractional).
  
            fx_train_or_state_0:
              either (a) output of the network at `t == 0` on the training set or (b)
              complete ODE state (`predict.ODEState`). Pass an ODE state if you want
              to operate on the full ODE state instead of output variables only
              (useful for inspecting auxiliary variables or resuming an optimizer with
              auxiliary variables from a specific state). Note that only
              `momentum != None` optimizer currently has auxiliary variables. To
              initialize an ODE state from scratch, call
              `predict.ODEState(fx_train_0, fx_test_0)`. If an ODE state is passed, an
              ODE state is returned. `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 predictions on the training set.
  
          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`.
            Alternatively can return an `ODEState` at time[s] `t`.
  
          Raises:
            ValueError: if `fx_test_0` is not `None`, but `k_test_train` is `None`.
          """
          _check_inputs(fx_train_or_state_0, fx_test_0, k_test_train)
  
          t = (
              jnp.array(t if t is not None else jnp.inf, dtype)
              * learning_rate
          )
          t_shape = t.shape
          t = t.reshape((-1,))
  
          # ODE solver requires `t[0]` to be the time when `fx_train_0` [and
          # `fx_test_0`] are evaluated, but also a strictly increasing sequence of
          # timesteps, so we always temporarily append an [almost] `0` at the start.
          t0 = jnp.where(
              t[0] == 0,
              jnp.full((1,), -1e-24, t.dtype),
              jnp.zeros((1,), t.dtype),
          )
          t = jnp.concatenate([t0, t])
  
          # Solve the ODE.
          fx_test_shape = _get_fx_test_shape(
              y_train, k_test_train, trace_axes
          )
          state_0 = get_state_0(
              fx_train_or_state_0, fx_test_0, fx_test_shape
          )
          state_t = ode.odeint(get_dstate_dt(k_test_train), state_0, t)
  
          # Remove the added `t0`.
          trim = lambda x: x[1:].reshape(t_shape + x.shape[1:])
          trim_tree = lambda tree: tree_map(trim, tree)
          state_t = trim_tree(state_t)
  
          # `ODEState` -> `ODEState`
          if isinstance(fx_train_or_state_0, ODEState):
              return state_t
  
          # `jnp.ndarray` -> `jnp.ndarray`
          fx_train_t, fx_test_t = state_t.fx_train, state_t.fx_test
  
          if fx_train_or_state_0 is not None and fx_test_0 is None:
              return fx_train_t
          if fx_test_0 is not None and fx_train_or_state_0 is None:
              return fx_test_t
          return fx_train_t, fx_test_t
  
      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 2 for calling a global function with example usage: gradient_descent(loss, k_train_train, y_train, learning_rate, momentum, trace_axes) and returns: predict_fn, ODEState, dstate_dt, fx_train_t,  fx_test_t, ODEState, state_t, fx_train_t, fx_test_t, ODEState