code: def pairwise_square_distance_tensor(
      x1, x2, feature_ndims, x1_example_ndims=1, x2_example_ndims=1
  ):
      """Returns pairwise distance between x1 and x2.
  
      This method is a generalization `pairwise_square_distance_matrix`.
      Given `x1` and `x2`, Tensors with shape `[..., N1, ... Nm, D1, ... Dk]` and
      `[..., M1, ... Ml, D1, ... Dk]`, compute the pairwise distance tensor `A` of
      shape `[..., N1, ... Nm, M1, ... Ml]`, where `m` is `x1_example_ndims` and
      `l` is `x2_example_ndims`.
  
      Args:
        x1: Floating point `Tensor` with shape `B1 + E1 + [D1, ..., Dk]`,
          where `B1` is a (possibly empty) batch shape, and `E1` is a list
          of `x1_example_ndims` values.
        x2: Floating point `Tensor` with shape `B2 + [M] + [D1, ..., Dk]`,
          where `B2` is a (possibly empty) batch shape that broadcasts
          with `B1`, and `E2` is a list of `x1_example_ndims` values.
        feature_ndims: The number of dimensions to consider for the euclidean
          norm. This is `k` from above.
        x1_example_ndims: Integer for number of example dimensions in `x1`. This is
          `len(E1)`.
        x2_example_ndims: Integer for number of example dimensions in `x2`. This is
          `len(E2)`.
      Returns:
        `Tensor` of shape `bc(B1, B2) + E1 + E2` representing the pairwise square
        distance tensor.
      """
      # Collapse all the example dimensions and then expand after.
      x1_shape = tf.shape(x1)
      x1_rank = ps.rank(x1)
      x1_example_shape = x1_shape[
          x1_rank
          - (feature_ndims + x1_example_ndims) : x1_rank
          - feature_ndims
      ]
  
      x2_shape = tf.shape(x2)
      x2_rank = ps.rank(x2)
      x2_example_shape = x2_shape[
          x2_rank
          - (feature_ndims + x2_example_ndims) : x2_rank
          - feature_ndims
      ]
  
      x1 = tf.reshape(
          x1,
          tf.concat(
              [
                  x1_shape[
                      : x1_rank - (feature_ndims + x1_example_ndims)
                  ],
                  [-1],
                  x1_shape[x1_rank - feature_ndims :],
              ],
              axis=0,
          ),
      )
      x2 = tf.reshape(
          x2,
          tf.concat(
              [
                  x2_shape[
                      : x2_rank - (feature_ndims + x2_example_ndims)
                  ],
                  [-1],
                  x2_shape[x2_rank - feature_ndims :],
              ],
              axis=0,
          ),
      )
      pairwise = pairwise_square_distance_matrix(
          x1, x2, feature_ndims=feature_ndims
      )
      # Now we need to undo the transformation.
      return tf.reshape(
          pairwise,
          tf.concat(
              [
                  tf.shape(pairwise)[:-2],
                  x1_example_shape,
                  x2_example_shape,
              ],
              axis=0,
          ),
      )
  
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 6 for calling a global function with example usage: pairwise_square_distance_tensor(x1, x2, feature_ndims, x1_example_ndims, x2_example_ndims) and returns: tf