code: def pairwise_squared_distance(x):
      """
      x : (n_samples, n_points, dims)
      return : (n_samples, n_points, n_points)
      """
      x2s = F.sum(x * x, -1, True)
      # assuming that __matmul__ is always implemented (true for PyTorch, MXNet and Chainer)
      return (
          x2s + F.swapaxes(x2s, -1, -2) - 2 * x @ F.swapaxes(x, -1, -2)
      )
  
imports: import copy
  from collections.abc import Iterable, Mapping
  import numpy
  import scipy.sparse
  import scipy.sparse.linalg
  from ..utils import version
  from ..None import backend, batch, convert, function, ndarray, subgraph, utils
  from .._ffi.function import _init_api
  from ..base import dgl_warning, DGLError, EID, NID
  from ..frame import Frame
  from ..heterograph import DGLGraph
  from ..heterograph_index import create_heterograph_from_relations, create_metagraph_index
  from ..partition import metis_partition, metis_partition_assignment, partition_graph_with_halo
  from ..sampling.neighbor import sample_neighbors
label_desc: Use image node_id 1 for calling a global function with example usage: pairwise_squared_distance(x) and returns: unknown