Change order of call args in RemoveAccidentalHits layer. (#45)

The order is now `logits`, `labels`, which is more consistent with the other layers and Keras in general.

Author: hertschuh

keras_rs/src/layers/retrieval/remove_accidental_hits.py

@@ -19,8 +19,8 @@ class RemoveAccidentalHits(keras.layers.Layer):
 
     def call(
         self,
-        labels: types.Tensor,
         logits: types.Tensor,
+        labels: types.Tensor,
         candidate_ids: types.Tensor,
     ) -> types.Tensor:
         """Zeroes selected logits.
@@ -29,16 +29,16 @@ def call(
         have the same ID as the positive candidate in that row.
 
         Args:
-          labels: one-hot labels tensor, typically
-              `[batch_size, num_candidates]` but can have more dimensions or be
-              1D as `[num_candidates]`.
-          logits: logits tensor. Must have the same shape as `labels`.
-          candidate_ids: candidate identifiers tensor, can be `[num_candidates]`
-              or `[batch_size, num_candidates]` or have more dimensions as long
-              as they match the last dimensions of `labels`.
+            logits: logits tensor, typically `[batch_size, num_candidates]` but
+                can have more dimensions or be 1D as `[num_candidates]`.
+            labels: one-hot labels tensor, must be the same shape as `logits`.
+            candidate_ids: candidate identifiers tensor, can be
+                `[num_candidates]` or `[batch_size, num_candidates]` or have
+                more dimensions as long as they match the last dimensions of
+                `labels`.
 
         Returns:
-          logits: Modified logits.
+            logits: Modified logits with the same shape as the input logits.
         """
         # A more principled way is to implement
         # `softmax_cross_entropy_with_logits` with a input mask. Here we

keras_rs/src/layers/retrieval/remove_accidental_hits_test.py

@@ -10,23 +10,24 @@
 
 class RemoveAccidentalHitsTest(testing.TestCase, parameterized.TestCase):
     def create_inputs(self, logits_rank=2, candidate_ids_rank=1):
-        num_candidates = 10
-        shape_3d = (15, 20, num_candidates)
+        shape_3d = (15, 20, 10)
         shape = shape_3d[-logits_rank:]
         candidate_ids_shape = shape_3d[-candidate_ids_rank:]
-        rng = keras.random.SeedGenerator(42)
+        num_candidates = shape[-1]
 
+        rng = keras.random.SeedGenerator(42)
+        logits = keras.random.uniform(shape, seed=rng)
         labels = keras.ops.one_hot(
             keras.random.randint(
                 shape[:-1], minval=0, maxval=num_candidates, seed=rng
             ),
             num_candidates,
         )
-        logits = keras.random.uniform(shape, seed=rng)
         candidate_ids = keras.random.randint(
             candidate_ids_shape, minval=0, maxval=num_candidates, seed=rng
         )
-        return labels, logits, candidate_ids
+
+        return logits, labels, candidate_ids
 
     @parameterized.named_parameters(
         [
@@ -63,12 +64,12 @@ def create_inputs(self, logits_rank=2, candidate_ids_rank=1):
         ]
     )
     def test_call(self, logits_rank, candidate_ids_rank):
-        labels, logits, candidate_ids = self.create_inputs(
+        logits, labels, candidate_ids = self.create_inputs(
             logits_rank=logits_rank, candidate_ids_rank=candidate_ids_rank
         )
 
         out_logits = remove_accidental_hits.RemoveAccidentalHits()(
-            labels, logits, candidate_ids
+            logits, labels, candidate_ids
         )
 
         # Logits of labels are unchanged.
@@ -148,36 +149,36 @@ def test_mismatched_labels_candidates_shapes(self):
 
     def test_predict(self):
         # Note: for predict, we test with probabilities that have a batch dim.
-        labels, logits, candidate_ids = self.create_inputs(candidate_ids_rank=2)
+        logits, labels, candidate_ids = self.create_inputs(candidate_ids_rank=2)
 
         layer = remove_accidental_hits.RemoveAccidentalHits()
-        in_labels = keras.layers.Input(labels.shape[1:])
         in_logits = keras.layers.Input(logits.shape[1:])
+        in_labels = keras.layers.Input(labels.shape[1:])
         in_candidate_ids = keras.layers.Input(labels.shape[1:])
-        out_logits = layer(in_labels, in_logits, in_candidate_ids)
+        out_logits = layer(in_logits, in_labels, in_candidate_ids)
         model = keras.Model(
-            [in_labels, in_logits, in_candidate_ids], out_logits
+            [in_logits, in_labels, in_candidate_ids], out_logits
         )
 
-        model.predict([labels, logits, candidate_ids], batch_size=8)
+        model.predict([logits, labels, candidate_ids], batch_size=8)
 
     def test_serialization(self):
         layer = remove_accidental_hits.RemoveAccidentalHits()
         restored = deserialize(serialize(layer))
         self.assertDictEqual(layer.get_config(), restored.get_config())
 
     def test_model_saving(self):
-        labels, logits, candidate_ids = self.create_inputs()
+        logits, labels, candidate_ids = self.create_inputs()
 
         layer = remove_accidental_hits.RemoveAccidentalHits()
-        in_labels = keras.layers.Input(labels.shape[1:])
         in_logits = keras.layers.Input(logits.shape[1:])
+        in_labels = keras.layers.Input(labels.shape[1:])
         in_candidate_ids = keras.layers.Input(batch_shape=candidate_ids.shape)
-        out_logits = layer(in_labels, in_logits, in_candidate_ids)
+        out_logits = layer(in_logits, in_labels, in_candidate_ids)
         model = keras.Model(
-            [in_labels, in_logits, in_candidate_ids], out_logits
+            [in_logits, in_labels, in_candidate_ids], out_logits
         )
 
         self.run_model_saving_test(
-            model=model, input_data=[labels, logits, candidate_ids]
+            model=model, input_data=[logits, labels, candidate_ids]
         )