code: def gram(kernel, xs):
      """Compute a Gram matrix from a kernel and an array of data points.
  
      Args:
        kernel: callable, maps pairs of data points to scalars.
        xs: array of data points, stacked along the leading dimension.
  
      Returns:
        A 2d array `a` such that `a[i, j] = kernel(xs[i], xs[j])`.
      """
      return vmap(lambda x: vmap(lambda y: kernel(x, y))(xs))(xs)
  
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 1 for calling a global function with example usage: gram(kernel, xs) and returns: vmap