Merge commit for internal changes

examples/BUILD

@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
+load("//third_party/bazel_rules/rules_python/python:py_binary.bzl", "py_binary")
+
 licenses(["notice"])
 
 package(

setup.py

@@ -27,7 +27,7 @@
 # This version number should always be that of the *next* (unreleased) version.
 # Immediately after uploading a package to PyPI, you should increment the
 # version number and push to gitHub.
-__version__ = "2.0.11"
+__version__ = "2.0.12"
 
 if "--release" in sys.argv:
   sys.argv.remove("--release")
@@ -45,7 +45,6 @@
     "scikit-learn",
     "matplotlib",
     "graphviz",
-    "dm-sonnet",
 ]
 
 # Part of the visualization code uses colabtools and IPython libraries. These

tensorflow_lattice/BUILD

@@ -28,14 +28,15 @@ py_library(
     srcs = [
         "__init__.py",
         "layers/__init__.py",
-        "sonnet_modules/__init__.py",
     ],
     srcs_version = "PY2AND3",
     deps = [
         "//tensorflow_lattice/python:aggregation_layer",
         "//tensorflow_lattice/python:categorical_calibration_layer",
         "//tensorflow_lattice/python:categorical_calibration_lib",
         "//tensorflow_lattice/python:cdf_layer",
+        "//tensorflow_lattice/python:conditional_cdf",
+        "//tensorflow_lattice/python:conditional_pwl_calibration",
         "//tensorflow_lattice/python:configs",
         "//tensorflow_lattice/python:estimators",
         "//tensorflow_lattice/python:kronecker_factored_lattice_layer",
@@ -50,7 +51,6 @@ py_library(
         "//tensorflow_lattice/python:premade_lib",
         "//tensorflow_lattice/python:pwl_calibration_layer",
         "//tensorflow_lattice/python:pwl_calibration_lib",
-        "//tensorflow_lattice/python:pwl_calibration_sonnet_module",
         "//tensorflow_lattice/python:rtl_layer",
         "//tensorflow_lattice/python:test_utils",
         "//tensorflow_lattice/python:utils",

tensorflow_lattice/__init__.py

@@ -20,11 +20,12 @@
 from __future__ import absolute_import
 
 import tensorflow_lattice.layers
-
 from tensorflow_lattice.python import aggregation_layer
 from tensorflow_lattice.python import categorical_calibration_layer
 from tensorflow_lattice.python import categorical_calibration_lib
 from tensorflow_lattice.python import cdf_layer
+from tensorflow_lattice.python import conditional_cdf
+from tensorflow_lattice.python import conditional_pwl_calibration
 from tensorflow_lattice.python import configs
 from tensorflow_lattice.python import estimators
 from tensorflow_lattice.python import kronecker_factored_lattice_layer
@@ -39,9 +40,6 @@
 from tensorflow_lattice.python import premade_lib
 from tensorflow_lattice.python import pwl_calibration_layer
 from tensorflow_lattice.python import pwl_calibration_lib
-from tensorflow_lattice.python import pwl_calibration_sonnet_module
 from tensorflow_lattice.python import test_utils
 from tensorflow_lattice.python import utils
 from tensorflow_lattice.python import visualization
-
-import tensorflow_lattice.sonnet_modules

tensorflow_lattice/python/BUILD

@@ -13,6 +13,9 @@
 # limitations under the License.
 # ==============================================================================
 
+load("//third_party/bazel_rules/rules_python/python:py_library.bzl", "py_library")
+load("//third_party/bazel_rules/rules_python/python:py_test.bzl", "py_test")
+
 package(
     default_visibility = [
         "//tensorflow_lattice:__subpackages__",
@@ -434,19 +437,6 @@ py_library(
     ],
 )
 
-py_library(
-    name = "pwl_calibration_sonnet_module",
-    srcs = ["pwl_calibration_sonnet_module.py"],
-    srcs_version = "PY2AND3",
-    deps = [
-        ":pwl_calibration_lib",
-        ":utils",
-        # absl/logging dep,
-        # sonnet dep,
-        # tensorflow:tensorflow_no_contrib dep,
-    ],
-)
-
 py_test(
     name = "pwl_calibration_test",
     size = "large",
@@ -457,14 +447,13 @@ py_test(
     deps = [
         ":parallel_combination_layer",
         ":pwl_calibration_layer",
-        ":pwl_calibration_sonnet_module",
         ":test_utils",
         ":utils",
         # absl/logging dep,
         # absl/testing:parameterized dep,
         # numpy dep,
-        # sonnet dep,
         # tensorflow dep,
+        # tensorflow:tensorflow_no_contrib dep,
     ],
 )
 
@@ -514,8 +503,6 @@ py_library(
         ":visualization",
         # absl/logging dep,
         # numpy dep,
-        # sonnet dep,
-        # tensorflow dep,
     ],
 )
 
@@ -528,6 +515,42 @@ py_library(
     ],
 )
 
+py_library(
+    name = "conditional_pwl_calibration",
+    srcs = ["conditional_pwl_calibration.py"],
+    deps = [
+        # numpy dep,
+        # tensorflow:tensorflow_no_contrib dep,
+    ],
+)
+
+py_library(
+    name = "conditional_cdf",
+    srcs = ["conditional_cdf.py"],
+    deps = [
+        # tensorflow:tensorflow_no_contrib dep,
+    ],
+)
+
+py_test(
+    name = "conditional_cdf_test",
+    srcs = ["conditional_cdf_test.py"],
+    deps = [
+        ":conditional_cdf",
+        # absl/testing:parameterized dep,
+        # tensorflow:tensorflow_no_contrib dep,
+    ],
+)
+
+py_test(
+    name = "conditional_pwl_calibration_test",
+    srcs = ["conditional_pwl_calibration_test.py"],
+    deps = [
+        ":conditional_pwl_calibration",
+        # tensorflow:tensorflow_no_contrib dep,
+    ],
+)
+
 py_test(
     name = "utils_test",
     srcs = ["utils_test.py"],

tensorflow_lattice/python/aggregation_layer.py

@@ -76,11 +76,14 @@ def call(self, x):
   def get_config(self):
     """Standard Keras get_config() method."""
     config = super(Aggregation, self).get_config().copy()
-    config.update({'model': tf.keras.utils.serialize_keras_object(self.model)})
+    config.update(
+        {'model': tf.keras.utils.legacy.serialize_keras_object(self.model)}
+    )
     return config
 
   @classmethod
   def from_config(cls, config, custom_objects=None):
-    model = tf.keras.utils.deserialize_keras_object(
-        config.pop('model'), custom_objects=custom_objects)
+    model = tf.keras.utils.legacy.deserialize_keras_object(
+        config.pop('model'), custom_objects=custom_objects
+    )
     return cls(model, **config)

tensorflow_lattice/python/categorical_calibration_layer.py

@@ -249,9 +249,11 @@ def get_config(self):
         "output_max": self.output_max,
         "monotonicities": self.monotonicities,
         "kernel_initializer":
-            keras.initializers.serialize(self.kernel_initializer),
+            keras.initializers.serialize(
+                self.kernel_initializer, use_legacy_format=True),
         "kernel_regularizer":
-            [keras.regularizers.serialize(r) for r in self.kernel_regularizer],
+            [keras.regularizers.serialize(r, use_legacy_format=True)
+             for r in self.kernel_regularizer],
         "default_input_value": self.default_input_value,
         "split_outputs": self.split_outputs,
     }  # pyformat: disable

tensorflow_lattice/python/cdf_layer.py

@@ -255,7 +255,8 @@ def get_config(self):
         "sparsity_factor":
             self.sparsity_factor,
         "kernel_initializer":
-            tf.keras.initializers.serialize(self.kernel_initializer),
+            tf.keras.initializers.serialize(
+                self.kernel_initializer, use_legacy_format=True),
     }
     config.update(super(CDF, self).get_config())
     return config