Fix test where shapes are not matched to keras.Input shapes.

PiperOrigin-RevId: 297461518

Author: omalleyt12

neural_structured_learning/keras/layers/pairwise_distance_test.py

@@ -89,12 +89,13 @@ def _distance_fn(x, y):
       return np.mean(np.square(x - y))
 
     # Common input.
-    sources = np.array([1., 1., 1., 1.])
-    targets = np.array([[4., 3., 2., 1.]])
+    sources = np.array([[1., 1., 1., 1.]])
+    targets = np.array([[[4., 3., 2., 1.]]])
     unweighted_distance = _distance_fn(sources, targets)
 
     def _make_symbolic_weights_model():
       """Makes a model where the weights are provided as input."""
+      # Shape doesn't include batch dimension.
       inputs = {
           'sources': tf.keras.Input(4),
           'targets': tf.keras.Input((1, 4)),
@@ -104,7 +105,7 @@ def _make_symbolic_weights_model():
       outputs = pairwise_distance_fn(**inputs)
       return tf.keras.Model(inputs=inputs, outputs=outputs)
 
-    weights = np.array([[2.]])
+    weights = np.array([[[2.]]])
     expected_distance = unweighted_distance * weights
     model = _make_symbolic_weights_model()
     self.assertNear(
@@ -117,6 +118,7 @@ def _make_symbolic_weights_model():
 
     def _make_fixed_weights_model(weights):
       """Makes a model where the weights are a static constant."""
+      # Shape doesn't include batch dimension.
       inputs = {
           'sources': tf.keras.Input(4),
           'targets': tf.keras.Input((1, 4)),