code: def _mesh_eval(func, x_limits, y_limits, params, num_ticks):
      # Evaluate func on a 2D grid defined by x_limits and y_limits.
      x = jnp.linspace(*x_limits, num=num_ticks)
      y = jnp.linspace(*y_limits, num=num_ticks)
      X, Y = jnp.meshgrid(x, y)
      xy_vec = jnp.stack([X.ravel(), Y.ravel()]).T
      zs = vmap(func, in_axes=(0, None))(xy_vec, params)
      return X, Y, zs.reshape(X.shape)
  
file_path: jax/examples/advi.py
imports: from functools import partial
  import matplotlib.pyplot
  from jax import jit, grad, vmap
  from jax import random
  from jax.example_libraries import optimizers
  import jax.numpy
  import jax.scipy.stats.norm
label_desc: Use image node_id 5 for calling a global function with example usage: _mesh_eval(func, x_limits, y_limits, params, num_ticks) and returns: X,  Y,  zs