Enable tfrs models to handle non-scalar regularization losses by appl…

…ying reduce_sum instead of sum on regularization losses. PiperOrigin-RevId: 504961825
tensorflow · Apr 24, 2023 · 23ad075 · 23ad075
1 parent 5557f09
commit 23ad075
Show file tree

Hide file tree

Showing 46 changed files with 54 additions and 49 deletions.
diff --git a/setup.py b/setup.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/__init__.py b/tensorflow_recommenders/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/examples/__init__.py b/tensorflow_recommenders/examples/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/examples/movielens.py b/tensorflow_recommenders/examples/movielens.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/examples/nbtool.py b/tensorflow_recommenders/examples/nbtool.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/__init__.py b/tensorflow_recommenders/experimental/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/layers/__init__.py b/tensorflow_recommenders/experimental/layers/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/layers/embedding/__init__.py b/tensorflow_recommenders/experimental/layers/embedding/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/layers/embedding/partial_tpu_embedding.py b/tensorflow_recommenders/experimental/layers/embedding/partial_tpu_embedding.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/layers/embedding/partial_tpu_embedding_test.py b/tensorflow_recommenders/experimental/layers/embedding/partial_tpu_embedding_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/models/__init__.py b/tensorflow_recommenders/experimental/models/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/models/ranking.py b/tensorflow_recommenders/experimental/models/ranking.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/models/ranking_test.py b/tensorflow_recommenders/experimental/models/ranking_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/optimizers/__init__.py b/tensorflow_recommenders/experimental/optimizers/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/optimizers/clippy_adagrad.py b/tensorflow_recommenders/experimental/optimizers/clippy_adagrad.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/optimizers/clippy_adagrad_test.py b/tensorflow_recommenders/experimental/optimizers/clippy_adagrad_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/optimizers/composite_optimizer.py b/tensorflow_recommenders/experimental/optimizers/composite_optimizer.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/experimental/optimizers/composite_optimizer_test.py b/tensorflow_recommenders/experimental/optimizers/composite_optimizer_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/__init__.py b/tensorflow_recommenders/layers/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/blocks.py b/tensorflow_recommenders/layers/blocks.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/embedding/__init__.py b/tensorflow_recommenders/layers/embedding/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/embedding/tpu_embedding_layer.py b/tensorflow_recommenders/layers/embedding/tpu_embedding_layer.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/embedding/tpu_embedding_layer_test.py b/tensorflow_recommenders/layers/embedding/tpu_embedding_layer_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/factorized_top_k.py b/tensorflow_recommenders/layers/factorized_top_k.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/factorized_top_k_test.py b/tensorflow_recommenders/layers/factorized_top_k_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/feature_interaction/__init__.py b/tensorflow_recommenders/layers/feature_interaction/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/feature_interaction/dcn.py b/tensorflow_recommenders/layers/feature_interaction/dcn.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/feature_interaction/dcn_test.py b/tensorflow_recommenders/layers/feature_interaction/dcn_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/feature_interaction/dot_interaction.py b/tensorflow_recommenders/layers/feature_interaction/dot_interaction.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/feature_interaction/dot_interaction_test.py b/tensorflow_recommenders/layers/feature_interaction/dot_interaction_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/loss.py b/tensorflow_recommenders/layers/loss.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/layers/loss_test.py b/tensorflow_recommenders/layers/loss_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/metrics/__init__.py b/tensorflow_recommenders/metrics/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/metrics/factorized_top_k.py b/tensorflow_recommenders/metrics/factorized_top_k.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/metrics/factorized_top_k_test.py b/tensorflow_recommenders/metrics/factorized_top_k_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/models/__init__.py b/tensorflow_recommenders/models/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/models/base.py b/tensorflow_recommenders/models/base.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -59,7 +59,8 @@ def compute_loss(self, inputs, training: bool = False) -> tf.Tensor:
     """
 
     raise NotImplementedError(
-        "Implementers must implement the `compute_loss` method.")
+        "Implementers must implement the `compute_loss` method."
+    )
 
   def train_step(self, inputs):
     """Custom train step using the `compute_loss` method."""
@@ -68,7 +69,9 @@ def train_step(self, inputs):
       loss = self.compute_loss(inputs, training=True)
 
       # Handle regularization losses as well.
-      regularization_loss = sum(self.losses)
+      regularization_loss = tf.reduce_sum(
+          [tf.reduce_sum(loss) for loss in self.losses]
+      )
 
       total_loss = loss + regularization_loss
 
@@ -88,7 +91,9 @@ def test_step(self, inputs):
     loss = self.compute_loss(inputs, training=False)
 
     # Handle regularization losses as well.
-    regularization_loss = sum(self.losses)
+    regularization_loss = tf.reduce_sum(
+        [tf.reduce_sum(loss) for loss in self.losses]
+    )
 
     total_loss = loss + regularization_loss
 

diff --git a/tensorflow_recommenders/models/base_test.py b/tensorflow_recommenders/models/base_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/tasks/__init__.py b/tensorflow_recommenders/tasks/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/tasks/base.py b/tensorflow_recommenders/tasks/base.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/tasks/ranking.py b/tensorflow_recommenders/tasks/ranking.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/tasks/ranking_test.py b/tensorflow_recommenders/tasks/ranking_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/tasks/retrieval.py b/tensorflow_recommenders/tasks/retrieval.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/tasks/retrieval_test.py b/tensorflow_recommenders/tasks/retrieval_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tensorflow_recommenders/types.py b/tensorflow_recommenders/types.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/tools/build_api_docs.py b/tools/build_api_docs.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Recommenders Authors.
+# Copyright 2023 The TensorFlow Recommenders Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.