Updates GDN/SignalConv to TF2 API and refactors reparameterizers.

From TF 2.5, we can rely on tf.Variables contained in Layer attributes (e.g.
tf.Modules) to show up recursively in self.weights/self.trainable_weights. We
can use this here to simplify reparameterization.

With this change, we don't rely on Layer.add_weight() any more. Instead, we just
assign a tf.Variable as an attribute of the Layer. Alternatively, we allow the
attribute to be a callable (e.g. for conditioning kernels on some other
computation), or a special callable 'Parameter' which implements
reparameterization. Parameter derives from tf.Module. As a result, the contained
reparameterized tf.Variable is associated with the Layer, and will be trained.

PiperOrigin-RevId: 355009145
Change-Id: I9936142d0ce999b3ebc8f0e8e50ef6cc0a0b3e20

Author: Johannes Ballé

BUILD

@@ -17,7 +17,7 @@ py_library(
         "//tensorflow_compression/python/entropy_models:universal",
         "//tensorflow_compression/python/layers:gdn",
         "//tensorflow_compression/python/layers:initializers",
-        "//tensorflow_compression/python/layers:parameterizers",
+        "//tensorflow_compression/python/layers:parameters",
         "//tensorflow_compression/python/layers:signal_conv",
         "//tensorflow_compression/python/layers:soft_round",
         "//tensorflow_compression/python/ops:math_ops",

tensorflow_compression/__init__.py

@@ -34,7 +34,7 @@
 
 from tensorflow_compression.python.layers.gdn import *
 from tensorflow_compression.python.layers.initializers import *
-from tensorflow_compression.python.layers.parameterizers import *
+from tensorflow_compression.python.layers.parameters import *
 from tensorflow_compression.python.layers.signal_conv import *
 from tensorflow_compression.python.layers.soft_round import *
 

tensorflow_compression/all_tests.py

@@ -31,7 +31,7 @@
 
 from tensorflow_compression.python.layers.gdn_test import *
 from tensorflow_compression.python.layers.initializers_test import *
-from tensorflow_compression.python.layers.parameterizers_test import *
+from tensorflow_compression.python.layers.parameters_test import *
 from tensorflow_compression.python.layers.signal_conv_test import *
 from tensorflow_compression.python.layers.soft_round_test import *
 

tensorflow_compression/python/layers/BUILD

@@ -8,14 +8,17 @@ py_library(
     name = "gdn",
     srcs = ["gdn.py"],
     srcs_version = "PY3",
-    deps = [":parameterizers"],
+    deps = [":parameters"],
 )
 
 py_test(
     name = "gdn_test",
     srcs = ["gdn_test.py"],
     python_version = "PY3",
-    deps = [":gdn"],
+    deps = [
+        ":gdn",
+        ":parameters",
+    ],
 )
 
 py_library(
@@ -24,47 +27,40 @@ py_library(
     srcs_version = "PY3",
 )
 
+py_test(
+    name = "initializers_test",
+    srcs = ["initializers_test.py"],
+    python_version = "PY3",
+    deps = [":initializers"],
+)
+
 py_library(
-    name = "parameterizers",
-    srcs = ["parameterizers.py"],
+    name = "parameters",
+    srcs = ["parameters.py"],
     srcs_version = "PY3",
     deps = [
         "//tensorflow_compression/python/ops:math_ops",
         "//tensorflow_compression/python/ops:spectral_ops",
     ],
 )
 
+py_test(
+    name = "parameters_test",
+    srcs = ["parameters_test.py"],
+    python_version = "PY3",
+    deps = [":parameters"],
+)
+
 py_library(
     name = "signal_conv",
     srcs = ["signal_conv.py"],
     srcs_version = "PY3",
     deps = [
-        ":parameterizers",
+        ":parameters",
         "//tensorflow_compression/python/ops:padding_ops",
     ],
 )
 
-py_library(
-    name = "soft_round",
-    srcs = ["soft_round.py"],
-    srcs_version = "PY3",
-    deps = ["//tensorflow_compression/python/ops:soft_round_ops"],
-)
-
-py_test(
-    name = "initializers_test",
-    srcs = ["initializers_test.py"],
-    python_version = "PY3",
-    deps = [":initializers"],
-)
-
-py_test(
-    name = "parameterizers_test",
-    srcs = ["parameterizers_test.py"],
-    python_version = "PY3",
-    deps = [":parameterizers"],
-)
-
 py_test(
     name = "signal_conv_test",
     timeout = "long",
@@ -73,16 +69,26 @@ py_test(
     shard_count = 3,
     deps = [
         ":initializers",
-        ":parameterizers",
+        ":parameters",
         ":signal_conv",
     ],
 )
 
+py_library(
+    name = "soft_round",
+    srcs = ["soft_round.py"],
+    srcs_version = "PY3",
+    deps = ["//tensorflow_compression/python/ops:soft_round_ops"],
+)
+
 py_test(
     name = "soft_round_test",
     srcs = ["soft_round_test.py"],
     python_version = "PY3",
-    deps = [":soft_round"],
+    deps = [
+        ":soft_round",
+        "//tensorflow_compression/python/ops:soft_round_ops",
+    ],
 )
 
 filegroup(