docs: def _scaled_square_dist(self, X, X2):
          """
          Returns ((X - X2ᵀ)/lengthscales)².
          Due to the implementation and floating-point imprecision, the
          result may actually be very slightly negative for entries very
          close to each other.
  
          This function can deal with leading dimensions in X and X2. 
          In the sample case, where X and X2 are both 2 dimensional, 
          for example, X is [N, D] and X2 is [M, D], then a tensor of shape 
          [N, M] is returned. If X is [N1, S1, D] and X2 is [N2, S2, D] 
          then the output will be [N1, S1, N2, S2].
          """
  
          X = X / self.lengthscales
  
          if X2 is None:
              Xs = tf.reduce_sum(tf.square(X), axis=-1, keepdims=True)
              dist = -2 * tf.matmul(X, X, transpose_b=True)
              dist += Xs + tf.matrix_transpose(Xs)
              return dist
  
          Xs = tf.reduce_sum(tf.square(X), axis=-1)
          X2 = X2 / self.lengthscales
          X2s = tf.reduce_sum(tf.square(X2), axis=-1)
          dist = -2 * tf.tensordot(X, X2, [[-1], [-1]])
          dist += _broadcasting_elementwise_op(tf.add, Xs, X2s)
          return dist