code: def train(kernel, xs, ys, regularization=0.01):
      gram_ = jit(partial(gram, kernel))
      gram_mat = gram_(xs)
      n = xs.shape[0]
  
      def objective(v):
          risk = 0.5 * jnp.sum((jnp.dot(gram_mat, v) - ys) ** 2.0)
          reg = regularization * jnp.sum(v**2.0)
          return risk + reg
  
      v = minimize(objective, jnp.zeros(n))
  
      def predict(x):
          prods = vmap(lambda x_: kernel(x, x_))(xs)
          return jnp.sum(v * prods)
  
      return jit(vmap(predict))
  
file_path: jax/examples/kernel_lsq.py
imports: from functools import partial
  import numpy.random
  import jax.numpy
  from jax.example_libraries import optimizers
  from jax import grad, jit, make_jaxpr, vmap, lax
label_desc: Use image node_id 3 for calling a global function with example usage: train(kernel, xs, ys, regularization) and returns: jit, unknown, jnp