Change the default value of keep_rank and flip its semantics when unpacking neighbor features. The point of reference is the (pre-batch) rank of tensors for the corresponding sample features.

PiperOrigin-RevId: 315433595

Author: arjung

neural_structured_learning/keras/layers/neighbor_features.py

@@ -99,7 +99,7 @@ def __init__(self,
     self._feature_names = (
         feature_names if feature_names is None else set(feature_names))
 
-  def call(self, inputs, keep_rank=False):
+  def call(self, inputs, keep_rank=True):
     """Extracts neighbor features and weights from a dictionary of inputs.
 
     This function is a wrapper around `utils.unpack_neighbor_features`. See
@@ -109,9 +109,10 @@ def call(self, inputs, keep_rank=False):
     Args:
       inputs: Dictionary of `tf.Tensor` features with keys for neighbors and
         weights described by `neighbor_config`.
-      keep_rank: Defaults to `False`. If `True`, each value of
-        `neighbor_features` will have an extra neighborhood size dimension at
-        axis 1.
+      keep_rank: Boolean indicating whether each value of `neighbor_features`
+        retains the rank from the corresponding value in `sample_features`
+        by merging the neighborhood size with the batch_size dimension, or
+        contains an extra neighborhood dimension at axis 1. Defaults to `True`.
 
     Returns:
       A tuple (sample_features, neighbor_features, neighbor_weights) of tensors.

neural_structured_learning/keras/layers/neighbor_features_test.py

@@ -35,7 +35,9 @@ def _make_model(neighbor_config, inputs, keep_rank, weight_dtype=None):
   Args:
     neighbor_config: An instance of `configs.GraphNeighborConfig`.
     inputs: A `tf.keras.Input` or a nested structure of `tf.keras.Input`s.
-    keep_rank: Whether to keep the extra neighborhood size dimention.
+    keep_rank: Whether to retain the rank of the original input tensors by
+      merging the neighborhood size with the batch_size dimension, or add an
+      extra neighborhood size dimension.
     weight_dtype: Optional `tf.DType` for weights.
 
   Returns:
@@ -105,15 +107,15 @@ def testDense(self, keep_rank):
     # Check that neighbors and weights are grouped together for each sample.
     for i in range(batch_size):
       self.assertAllEqual(
-          neighbors[i] if keep_rank else
-          neighbors[(i * num_neighbors):((i + 1) * num_neighbors)],
+          neighbors[(i * num_neighbors):((i + 1) * num_neighbors)]
+          if keep_rank else neighbors[i],
           np.stack([
               features['NL_nbr_0_image'][i],
               features['NL_nbr_1_image'][i],
               features['NL_nbr_2_image'][i],
           ]))
       self.assertAllEqual(
-          weights[i] if keep_rank else np.split(weights, batch_size)[i],
+          np.split(weights, batch_size)[i] if keep_rank else weights[i],
           np.stack([
               features['NL_nbr_0_weight'][i],
               features['NL_nbr_1_weight'][i],
@@ -160,8 +162,8 @@ def testSparse(self, keep_rank):
     self.assertAllClose(
         weights,
         np.array([0.9, 0.25, 0.3, 0., 0.6, 0.75, 0.,
-                  0.]).reshape((batch_size, 2,
-                                1) if keep_rank else (batch_size * 2, 1)))
+                  0.]).reshape((batch_size * 2,
+                                1) if keep_rank else (batch_size, 2, 1)))
     # Check that neighbors are grouped together.
     dense_neighbors = self.evaluate(tf.sparse.to_dense(neighbors['input'], -1.))
     neighbor0 = self.evaluate(
@@ -170,8 +172,8 @@ def testSparse(self, keep_rank):
         tf.sparse.to_dense(features['NL_nbr_1_input'], -1))
     for i in range(batch_size):
       actual = (
-          dense_neighbors[i]
-          if keep_rank else np.split(dense_neighbors, batch_size)[i])
+          np.split(dense_neighbors, batch_size)[i]
+          if keep_rank else dense_neighbors[i])
       self.assertAllEqual(actual, np.stack([neighbor0[i], neighbor1[i]]))
 
 

neural_structured_learning/lib/utils.py

@@ -443,15 +443,15 @@ def _interleave_and_merge(tensors,
     # This is the equivalent of tf.stack() for sparse tensors.
     concatenated_tensors = tf.sparse.concat(
         axis=1, sp_inputs=[tf.sparse.expand_dims(t, 1) for t in tensors])
-    return (concatenated_tensors if keep_rank else tf.sparse.reshape(
-        concatenated_tensors, shape=merged_shape))
+    return (tf.sparse.reshape(concatenated_tensors, shape=merged_shape)
+            if keep_rank else concatenated_tensors)
   else:
     stacked_tensors = tf.stack(tensors, axis=1)
-    return (stacked_tensors if keep_rank else tf.reshape(
-        stacked_tensors, shape=merged_shape))
+    return (tf.reshape(stacked_tensors, shape=merged_shape)
+            if keep_rank else stacked_tensors)
 
 
-def unpack_neighbor_features(features, neighbor_config, keep_rank=False):
+def unpack_neighbor_features(features, neighbor_config, keep_rank=True):
   """Extracts sample features, neighbor features, and neighbor weights.
 
   For example, suppose `features` contains a single sample feature named
@@ -520,8 +520,9 @@ def unpack_neighbor_features(features, neighbor_config, keep_rank=False):
       1]`, where `B` is the batch size. Neighbor weight tensors cannot be sparse
       tensors.
     neighbor_config: An instance of `nsl.configs.GraphNeighborConfig`.
-    keep_rank: Whether to preserve the neighborhood size dimension. Defaults to
-      `False`.
+    keep_rank: Boolean indicating whether to retain the rank from the input or
+      to introduce a new dimension for the neighborhood size (axis 1). Defaults
+      to `True`.
 
   Returns:
     sample_features: a dictionary mapping feature names to tensors. The shape