code: def _cdist(a: Array, b: Array, metric: MetricFn) -> Array:
      """Returns the distance between each pair of the two collections of inputs."""
      return jax.vmap(jax.vmap(metric, (None, 0)), (0, None))(a, b)
  
file_path: rlax/rlax/_src/episodic_memory.py
imports: import functools
  from typing import Callable
  import chex
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling a global function with example usage: _cdist(a, b, metric) and returns: jax