docs: def _cust_square_dist(A, B, sc):
          """
          Custom version of _square_dist that allows sc to provide per-datapoint length
          scales. sc: N x M x D.
          """
          return tf.reduce_sum(tf.square((tf.expand_dims(A, 1) - tf.expand_dims(B, 0)) / sc), 2)