code: def pairwise_square_distance_matrix(x1, x2, feature_ndims):
      """Returns pairwise square distance between x1 and x2.
  
      Given `x1` and `x2`, Tensors with shape `[..., N, D1, ... Dk]` and
      `[..., M, D1, ... Dk]`, compute the pairwise distance matrix `a_ij` of shape
      `[..., N, M]`, where each entry `a_ij` is the square of the euclidean norm of
      `x1[..., i, ...] - x2[..., j, ...]`.
  
      The approach uses the fact that (where k = 1).
      ```none
        a_ij = sum_d (x1[i, d] - x2[j, d]) ** 2 =
        sum_d x1[i, d] ** 2 + x2[j, d] ** 2 - 2 * x1[i, d] * x2[j, d]
      ```
  
      The latter term can be written as a matmul between `x1` and `x2`.
      This reduces the memory from the naive approach of computing the
      squared difference of `x1` and `x2` by a factor of `prod_k D_k`.
      This is at the cost of the computation being more numerically unstable.
  
      Args:
        x1: Floating point `Tensor` with shape `B1 + [N] + [D1, ..., Dk]`,
          where `B1` is a (possibly empty) batch shape.
        x2: Floating point `Tensor` with shape `B2 + [M] + [D1, ..., Dk]`,
          where `B2` is a (possibly empty) batch shape that broadcasts
          with `B1`.
        feature_ndims: The number of dimensions to consider for the euclidean
          norm. This is `k` from above.
      Returns:
        `Tensor` of shape `[..., N, M]` representing the pairwise square
        distance matrix.
      """
      row_norm_x1 = sum_rightmost_ndims_preserving_shape(
          tf.square(x1), feature_ndims
      )[..., tf.newaxis]
      row_norm_x2 = sum_rightmost_ndims_preserving_shape(
          tf.square(x2), feature_ndims
      )[..., tf.newaxis, :]
  
      x1_rank = ps.rank(x1)
      x2_rank = ps.rank(x2)
  
      reshaped_x1 = tf.reshape(
          x1,
          ps.concat(
              [
                  ps.shape(x1)[: x1_rank - feature_ndims],
                  [
                      ps.reduce_prod(
                          ps.shape(x1)[x1_rank - feature_ndims :]
                      )
                  ],
              ],
              axis=0,
          ),
      )
      reshaped_x2 = tf.reshape(
          x2,
          ps.concat(
              [
                  ps.shape(x2)[: x2_rank - feature_ndims],
                  [
                      ps.reduce_prod(
                          ps.shape(x2)[x2_rank - feature_ndims :]
                      )
                  ],
              ],
              axis=0,
          ),
      )
      pairwise_sq = (
          row_norm_x1
          + row_norm_x2
          - 2
          * tf.linalg.matmul(reshaped_x1, reshaped_x2, transpose_b=True)
      )
      return tf.clip_by_value(pairwise_sq, 0.0, np.inf)
  
file_path: probability/tensorflow_probability/python/math/psd_kernels/internal/util.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 5 for calling a global function with example usage: pairwise_square_distance_matrix(x1, x2, feature_ndims) and returns: tf