code: def _pairwise_squared_distance_matrix(x: TensorLike) -> tf.Tensor:
      """Pairwise squared distance among a (batch) matrix's rows (2nd dim).
  
      This saves a bit of computation vs. using
      `_cross_squared_distance_matrix(x, x)`
  
      Args:
        x: `[batch_size, n, d]` float `Tensor`.
  
      Returns:
        squared_dists: `[batch_size, n, n]` float `Tensor`, where
        `squared_dists[b,i,j] = ||x[b,i,:] - x[b,j,:]||^2`.
      """
  
      x_x_transpose = tf.matmul(x, x, adjoint_b=True)
      x_norm_squared = tf.linalg.diag_part(x_x_transpose)
      x_norm_squared_tile = tf.expand_dims(x_norm_squared, 2)
  
      # squared_dists[b,i,j] = ||x_bi - x_bj||^2 =
      # = x_bi'x_bi- 2x_bi'x_bj + x_bj'x_bj
      squared_dists = (
          x_norm_squared_tile
          - 2 * x_x_transpose
          + tf.transpose(x_norm_squared_tile, [0, 2, 1])
      )
  
      return squared_dists
  
file_path: addons/tensorflow_addons/image/interpolate_spline.py
imports: import tensorflow
  from tensorflow_addons.utils.types import FloatTensorLike, TensorLike
label_desc: Use image node_id 2 for calling a global function with example usage: _pairwise_squared_distance_matrix(x) and returns: squared_dists