code: def pairwise_squared_distance(x):
      """
      x : (n_samples, n_points, dims)
      return : (n_samples, n_points, n_points)
      """
      x2s = (x * x).sum(-1, keepdim=True)
      return x2s + x2s.transpose(-1, -2) - 2 * x @ x.transpose(-1, -2)
  
file_path: dgl/python/dgl/nn/pytorch/factory.py
imports: import torch.nn
  from ...transforms import knn_graph, radius_graph, segmented_knn_graph
label_desc: Use image node_id 1 for calling a global function with example usage: pairwise_squared_distance(x) and returns: unknown