code: def mse(a: jnp.ndarray, b: jnp.ndarray) -> jnp.ndarray:
      """MSE distance."""
      return jnp.mean(jnp.square(a - b))
  
file_path: acme/acme/agents/jax/mbop/losses.py
imports: import dataclasses
  from typing import Any, Callable, Optional, Tuple, Union
  from acme import types
  from acme.agents.jax.mbop import dataset
  from acme.jax import networks
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: mse(a, b) and returns: jnp