def squared_euclidean_distance(a: jnp.ndarray,  # pytype: disable=annotation-type-mismatch  # jax-ndarray
                               b: jnp.ndarray,
                               b2: jnp.ndarray = None,
                               precision=DEFAULT_PRECISION) -> jnp.ndarray:
  """Computes the pairwise squared Euclidean distance.

  Args:
    a: float32: (n, d): An array of points.
    b: float32: (m, d): An array of points.
    b2: float32: (d, m): b square transpose.
    precision: use HIGHEST precision by default

  Returns:
    d: float32: (n, m): Where d[i, j] is the squared Euclidean distance between
    a[i] and b[j].
  """
  if b2 is None:
    b2 = jnp.sum(b.T**2, axis=0, keepdims=True)
  a2 = jnp.sum(a**2, axis=1, keepdims=True)
  ab = jnp.matmul(a, b.T, precision=precision)
  d = a2 - 2 * ab + b2
  return d
Computes the pairwise squared Euclidean distance.

  Args:
    a: float32: (n, d): An array of points.
    b: float32: (m, d): An array of points.
    b2: float32: (d, m): b square transpose.
    precision: use HIGHEST precision by default

  Returns:
    d: float32: (n, m): Where d[i, j] is the squared Euclidean distance between
    a[i] and b[j].
  
videogvt/train_lib/losses.py