code: def _sqeuclidian(x: Array, y: Array) -> Array:
      return jnp.sum(jnp.square(x - y))
  
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 1 for calling a global function with example usage: _sqeuclidian(x, y) and returns: jnp