code: def _pairwise_distances(coordinates):
      xx = jnp.sum(coordinates**2, axis=-1)
      xy = jnp.einsum("...xy,...zy->...xz", coordinates, coordinates)
  
      x2 = jnp.expand_dims(xx, axis=1)
      y2 = jnp.expand_dims(xx, axis=1)
      sq_distances = x2 + y2 - 2 * xy
      return sq_distances
  
imports: from typing import Any
  import flax.linen
  from flax.linen.initializers import zeros
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.projects.performer import performer
label_desc: Use image node_id 1 for calling a global function with example usage: _pairwise_distances(coordinates) and returns: sq_distances