Convert PSDKernels to AutoCompositeTensor.

emilyfertig · tensorflower-gardener · commit f58370f03f69 · 2021-05-24T11:25:11.000-07:00
PiperOrigin-RevId: 375517074
diff --git a/tensorflow_probability/python/distributions/variational_gaussian_process.py b/tensorflow_probability/python/distributions/variational_gaussian_process.py
@@ -35,15 +35,16 @@
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
 from tensorflow_probability.python.internal import tensorshape_util
-from tensorflow_probability.python.math import psd_kernels as tfpk
+from tensorflow_probability.python.math.psd_kernels import positive_semidefinite_kernel as psd_kernel
 from tensorflow_probability.python.math.psd_kernels.internal import util as kernel_util
 
 __all__ = [
     'VariationalGaussianProcess',
 ]
 
 
-class _VariationalKernel(tfpk.PositiveSemidefiniteKernel):
+@psd_kernel.auto_composite_tensor_psd_kernel
+class _VariationalKernel(psd_kernel.AutoCompositeTensorPsdKernel):
   """A PSDKernel which computes the variational kernel from [Titsias, 2009].
 
   The VariationalGaussianProcess can be cast as a special case of
@@ -129,12 +130,12 @@ def __init__(self,
           inducing_index_points, dtype=dtype, name='inducing_index_points')
       self._variational_scale = tensor_util.convert_nonref_to_tensor(
           variational_scale, dtype=dtype, name='variational_scale')
-      jitter = tensor_util.convert_nonref_to_tensor(
+      self._jitter = tensor_util.convert_nonref_to_tensor(
           jitter, dtype=dtype, name='jitter')
 
       def _compute_chol_kzz(z):
         kzz = base_kernel.matrix(z, z)
-        result = tf.linalg.cholesky(_add_diagonal_shift(kzz, jitter))
+        result = tf.linalg.cholesky(_add_diagonal_shift(kzz, self._jitter))
         return result
 
       # Somewhat confusingly, but for the sake of brevity, we use `var` to refer
diff --git a/tensorflow_probability/python/math/psd_kernels/BUILD b/tensorflow_probability/python/math/psd_kernels/BUILD
@@ -55,6 +55,7 @@ multi_substrate_py_library(
     deps = [
         # six dep,
         # tensorflow dep,
+        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:tensorshape_util",
     ],
@@ -66,9 +67,11 @@ multi_substrate_py_test(
     srcs = ["positive_semidefinite_kernel_test.py"],
     numpy_tags = ["notap"],
     deps = [
+        ":positive_semidefinite_kernel",
         # absl/testing:parameterized dep,
         # tensorflow dep,
         "//tensorflow_probability",
+        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:test_util",
         "//tensorflow_probability/python/math/psd_kernels/internal:util",
     ],
diff --git a/tensorflow_probability/python/math/psd_kernels/exp_sin_squared.py b/tensorflow_probability/python/math/psd_kernels/exp_sin_squared.py
@@ -23,13 +23,13 @@
 
 from tensorflow_probability.python.internal import assert_util
 from tensorflow_probability.python.internal import tensor_util
+from tensorflow_probability.python.math.psd_kernels import positive_semidefinite_kernel as psd_kernel
 from tensorflow_probability.python.math.psd_kernels.internal import util
-from tensorflow_probability.python.math.psd_kernels.positive_semidefinite_kernel import PositiveSemidefiniteKernel
-
 __all__ = ['ExpSinSquared']
 
 
-class ExpSinSquared(PositiveSemidefiniteKernel):
+@psd_kernel.auto_composite_tensor_psd_kernel
+class ExpSinSquared(psd_kernel.AutoCompositeTensorPsdKernel):
   """Exponentiated Sine Squared Kernel.
 
     Also known as the "Periodic" kernel, this kernel, when
diff --git a/tensorflow_probability/python/math/psd_kernels/exponentiated_quadratic.py b/tensorflow_probability/python/math/psd_kernels/exponentiated_quadratic.py
@@ -22,14 +22,15 @@
 
 from tensorflow_probability.python.internal import assert_util
 from tensorflow_probability.python.internal import tensor_util
+from tensorflow_probability.python.math.psd_kernels import positive_semidefinite_kernel as psd_kernel
 from tensorflow_probability.python.math.psd_kernels.internal import util
-from tensorflow_probability.python.math.psd_kernels.positive_semidefinite_kernel import PositiveSemidefiniteKernel
 
 
 __all__ = ['ExponentiatedQuadratic']
 
 
-class ExponentiatedQuadratic(PositiveSemidefiniteKernel):
+@psd_kernel.auto_composite_tensor_psd_kernel
+class ExponentiatedQuadratic(psd_kernel.AutoCompositeTensorPsdKernel):
   """The ExponentiatedQuadratic kernel.
 
   Sometimes called the "squared exponential", "Gaussian" or "radial basis
diff --git a/tensorflow_probability/python/math/psd_kernels/feature_scaled.py b/tensorflow_probability/python/math/psd_kernels/feature_scaled.py
@@ -21,15 +21,17 @@
 import tensorflow.compat.v2 as tf
 from tensorflow_probability.python.internal import assert_util
 from tensorflow_probability.python.internal import tensor_util
-from tensorflow_probability.python.math.psd_kernels.feature_transformed import FeatureTransformed
+from tensorflow_probability.python.math.psd_kernels import feature_transformed
+from tensorflow_probability.python.math.psd_kernels import positive_semidefinite_kernel as psd_kernel
 from tensorflow_probability.python.math.psd_kernels.internal import util
 
 __all__ = ['FeatureScaled']
 
 
 # TODO(b/132103412): Support more general scaling via LinearOperator, along with
 # scaling all feature dimensions.
-class FeatureScaled(FeatureTransformed):
+@psd_kernel.auto_composite_tensor_psd_kernel
+class FeatureScaled(feature_transformed.FeatureTransformed):
   """Kernel that first rescales all feature dimensions.
 
   Given a kernel `k` and `scale_diag` and inputs `x` and `y`, this kernel
diff --git a/tensorflow_probability/python/math/psd_kernels/feature_transformed.py b/tensorflow_probability/python/math/psd_kernels/feature_transformed.py
@@ -20,12 +20,13 @@
 
 import tensorflow.compat.v2 as tf
 
-from tensorflow_probability.python.math.psd_kernels.positive_semidefinite_kernel import PositiveSemidefiniteKernel
+from tensorflow_probability.python.math.psd_kernels import positive_semidefinite_kernel as psd_kernel
 
 __all__ = ['FeatureTransformed']
 
 
-class FeatureTransformed(PositiveSemidefiniteKernel):
+@psd_kernel.auto_composite_tensor_psd_kernel
+class FeatureTransformed(psd_kernel.AutoCompositeTensorPsdKernel):
   """Input transformed kernel.
 
   Given a kernel `k` and function `f`, compute `k_{new}(x, y) = k(f(x), f(y))`.
diff --git a/tensorflow_probability/python/math/psd_kernels/kumaraswamy_transformed.py b/tensorflow_probability/python/math/psd_kernels/kumaraswamy_transformed.py
@@ -25,12 +25,14 @@
 from tensorflow_probability.python.internal import assert_util
 from tensorflow_probability.python.internal import tensor_util
 from tensorflow_probability.python.math.psd_kernels import feature_transformed
+from tensorflow_probability.python.math.psd_kernels import positive_semidefinite_kernel as psd_kernel
 from tensorflow_probability.python.math.psd_kernels.internal import util
 
 
 __all__ = ['KumaraswamyTransformed']
 
 
+@psd_kernel.auto_composite_tensor_psd_kernel
 class KumaraswamyTransformed(feature_transformed.FeatureTransformed):
   """Transform inputs by Kumaraswamy bijector.
 
diff --git a/tensorflow_probability/python/math/psd_kernels/matern.py b/tensorflow_probability/python/math/psd_kernels/matern.py
@@ -25,8 +25,8 @@
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import tensor_util
 from tensorflow_probability.python.math import bessel as tfp_math
+from tensorflow_probability.python.math.psd_kernels import positive_semidefinite_kernel as psd_kernel
 from tensorflow_probability.python.math.psd_kernels.internal import util
-from tensorflow_probability.python.math.psd_kernels.positive_semidefinite_kernel import PositiveSemidefiniteKernel
 
 __all__ = [
     'GeneralizedMatern',
@@ -111,7 +111,9 @@ def _tensor(self, x1, x2, x1_example_ndims, x2_example_ndims):
         example_ndims=(x1_example_ndims + x2_example_ndims))
 
 
-class GeneralizedMatern(_AmplitudeLengthScaleMixin, PositiveSemidefiniteKernel):
+@psd_kernel.auto_composite_tensor_psd_kernel
+class GeneralizedMatern(_AmplitudeLengthScaleMixin,
+                        psd_kernel.AutoCompositeTensorPsdKernel):
   """Generalized Matern Kernel.
 
   This kernel parameterizes the Matern family of kernels.
@@ -224,7 +226,9 @@ def _parameter_control_dependencies(self, is_init):
     return assertions
 
 
-class MaternOneHalf(_AmplitudeLengthScaleMixin, PositiveSemidefiniteKernel):
+@psd_kernel.auto_composite_tensor_psd_kernel
+class MaternOneHalf(_AmplitudeLengthScaleMixin,
+                    psd_kernel.AutoCompositeTensorPsdKernel):
   """Matern Kernel with parameter 1/2.
 
   This kernel is part of the Matern family of kernels, with parameter 1/2.
@@ -295,7 +299,9 @@ def _apply_with_distance(
     return tf.exp(log_result)
 
 
-class MaternThreeHalves(_AmplitudeLengthScaleMixin, PositiveSemidefiniteKernel):
+@psd_kernel.auto_composite_tensor_psd_kernel
+class MaternThreeHalves(_AmplitudeLengthScaleMixin,
+                        psd_kernel.AutoCompositeTensorPsdKernel):
   """Matern Kernel with parameter 3/2.
 
   This kernel is part of the Matern family of kernels, with parameter 3/2.
@@ -363,7 +369,9 @@ def _apply_with_distance(
     return tf.exp(log_result)
 
 
-class MaternFiveHalves(_AmplitudeLengthScaleMixin, PositiveSemidefiniteKernel):
+@psd_kernel.auto_composite_tensor_psd_kernel
+class MaternFiveHalves(_AmplitudeLengthScaleMixin,
+                       psd_kernel.AutoCompositeTensorPsdKernel):
   """Matern 5/2 Kernel.
 
   This kernel is part of the Matern family of kernels, with parameter 5/2.
diff --git a/tensorflow_probability/python/math/psd_kernels/parabolic.py b/tensorflow_probability/python/math/psd_kernels/parabolic.py
@@ -23,13 +23,14 @@
 from tensorflow_probability.python.internal import assert_util
 from tensorflow_probability.python.internal import prefer_static as ps
 from tensorflow_probability.python.internal import tensor_util
+from tensorflow_probability.python.math.psd_kernels import positive_semidefinite_kernel as psd_kernel
 from tensorflow_probability.python.math.psd_kernels.internal import util
-from tensorflow_probability.python.math.psd_kernels.positive_semidefinite_kernel import PositiveSemidefiniteKernel
 
 __all__ = ['Parabolic']
 
 
-class Parabolic(PositiveSemidefiniteKernel):
+@psd_kernel.auto_composite_tensor_psd_kernel
+class Parabolic(psd_kernel.AutoCompositeTensorPsdKernel):
   """The Parabolic kernel.
 
   ```none
diff --git a/tensorflow_probability/python/math/psd_kernels/polynomial.py b/tensorflow_probability/python/math/psd_kernels/polynomial.py
@@ -24,8 +24,8 @@
 
 from tensorflow_probability.python.internal import assert_util
 from tensorflow_probability.python.internal import tensor_util
+from tensorflow_probability.python.math.psd_kernels import positive_semidefinite_kernel as psd_kernel
 from tensorflow_probability.python.math.psd_kernels.internal import util
-from tensorflow_probability.python.math.psd_kernels.positive_semidefinite_kernel import PositiveSemidefiniteKernel
 
 __all__ = [
     'Linear',
@@ -45,7 +45,8 @@ def _maybe_shape_dynamic(tensor):
   return tf.shape(tensor)
 
 
-class Polynomial(PositiveSemidefiniteKernel):
+@psd_kernel.auto_composite_tensor_psd_kernel
+class Polynomial(psd_kernel.AutoCompositeTensorPsdKernel):
   """Polynomial Kernel.
 
     Is based on the dot product covariance function and can be obtained
@@ -213,6 +214,7 @@ def _parameter_control_dependencies(self, is_init):
     return assertions
 
 
+@psd_kernel.auto_composite_tensor_psd_kernel
 class Linear(Polynomial):
   """Linear Kernel.
 
diff --git a/tensorflow_probability/python/math/psd_kernels/positive_semidefinite_kernel.py b/tensorflow_probability/python/math/psd_kernels/positive_semidefinite_kernel.py
@@ -25,6 +25,7 @@
 import six
 import tensorflow.compat.v2 as tf
 
+from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import tensorshape_util
 from tensorflow_probability.python.math.psd_kernels.internal import util
@@ -925,6 +926,18 @@ def _parameter_control_dependencies(self, is_init):
     return ()
 
 
+class AutoCompositeTensorPsdKernel(PositiveSemidefiniteKernel,
+                                   auto_composite_tensor.AutoCompositeTensor):
+  pass
+
+
+auto_composite_tensor_psd_kernel = functools.partial(
+    auto_composite_tensor.auto_composite_tensor,
+    omit_kwargs=('parameters',),
+    non_identifying_kwargs=('name',),
+    module_name='tfp.math._psdkernels')
+
+
 def _flatten_summand_list(kernels):
   """Flatten a list of kernels which may contain _SumKernel instances.
 
@@ -937,7 +950,7 @@ def _flatten_summand_list(kernels):
   """
   flattened = []
   for k in kernels:
-    if isinstance(k, _SumKernel):
+    if isinstance(k, _NonCompositeTensorSumKernel):
       flattened += k.kernels
     else:
       flattened.append(k)
@@ -956,14 +969,14 @@ def _flatten_multiplicand_list(kernels):
   """
   flattened = []
   for k in kernels:
-    if isinstance(k, _ProductKernel):
+    if isinstance(k, _NonCompositeTensorProductKernel):
       flattened += k.kernels
     else:
       flattened.append(k)
   return flattened
 
 
-class _SumKernel(PositiveSemidefiniteKernel):
+class _NonCompositeTensorSumKernel(PositiveSemidefiniteKernel):
   """Kernel class representing summation over a list of kernels.
 
   Mathematically this class represents the pointwise sum of several kernels.
@@ -1019,7 +1032,7 @@ def __init__(self, kernels, name=None):
     if name is None:
       name = 'SumKernel'
     # We have ensured the list is non-empty and all feature_ndims are the same.
-    super(_SumKernel, self).__init__(
+    super(_NonCompositeTensorSumKernel, self).__init__(
         feature_ndims=kernels[0].feature_ndims,
         dtype=util.maybe_get_common_dtype(
             [None if k.dtype is None else k for k in kernels]),
@@ -1052,7 +1065,25 @@ def _batch_shape_tensor(self):
                             [k.batch_shape_tensor() for k in self.kernels])
 
 
-class _ProductKernel(PositiveSemidefiniteKernel):
+@auto_composite_tensor_psd_kernel
+class _SumKernel(_NonCompositeTensorSumKernel,
+                 auto_composite_tensor.AutoCompositeTensor):
+
+  def __new__(cls, kernels, name=None):
+    if cls is _SumKernel:
+      if not all(isinstance(k, tf.__internal__.CompositeTensor)
+                 for k in kernels):
+        return _NonCompositeTensorSumKernel(kernels, name=name)
+    return super(_SumKernel, cls).__new__(cls)
+
+
+_SumKernel.__doc__ = _NonCompositeTensorSumKernel.__doc__ + '\n' + (
+    'When an `_SumKernel` is constructed, if any element of its `kernels`'
+    'list is not a `CompositeTensor` instance, a `_NonCompositeTensorSumKernel`'
+    ' instance is returned instead.')
+
+
+class _NonCompositeTensorProductKernel(PositiveSemidefiniteKernel):
   """Kernel class representing the product over a list of kernels.
 
   Mathematically this class represents the pointwise product of several kernels.
@@ -1098,7 +1129,7 @@ def __init__(self, kernels, name=None):
     if name is None:
       name = 'ProductKernel'
     # We have ensured the list is non-empty and all feature_ndims are the same.
-    super(_ProductKernel, self).__init__(
+    super(_NonCompositeTensorProductKernel, self).__init__(
         feature_ndims=kernels[0].feature_ndims,
         dtype=util.maybe_get_common_dtype(
             [None if k.dtype is None else k for k in kernels]),
@@ -1133,3 +1164,21 @@ def _batch_shape(self):
   def _batch_shape_tensor(self):
     return functools.reduce(tf.broadcast_dynamic_shape,
                             [k.batch_shape_tensor() for k in self.kernels])
+
+
+@auto_composite_tensor_psd_kernel
+class _ProductKernel(_NonCompositeTensorProductKernel,
+                     auto_composite_tensor.AutoCompositeTensor):
+
+  def __new__(cls, kernels, name=None):
+    if cls is _ProductKernel:
+      if not all(isinstance(k, tf.__internal__.CompositeTensor)
+                 for k in kernels):
+        return _NonCompositeTensorProductKernel(kernels, name=name)
+    return super(_ProductKernel, cls).__new__(cls)
+
+
+_ProductKernel.__doc__ = _NonCompositeTensorProductKernel.__doc__ + '\n' + (
+    'When an `_ProductKernel` is constructed, if any element of its `kernels`'
+    'list is not a `CompositeTensor` instance, a '
+    '`_NonCompositeTensorProductKernel` instance is returned instead.')
diff --git a/tensorflow_probability/python/math/psd_kernels/positive_semidefinite_kernel_test.py b/tensorflow_probability/python/math/psd_kernels/positive_semidefinite_kernel_test.py
diff --git a/tensorflow_probability/python/math/psd_kernels/psd_kernel_properties_test.py b/tensorflow_probability/python/math/psd_kernels/psd_kernel_properties_test.py
diff --git a/tensorflow_probability/python/math/psd_kernels/rational_quadratic.py b/tensorflow_probability/python/math/psd_kernels/rational_quadratic.py
diff --git a/tensorflow_probability/python/math/psd_kernels/schur_complement.py b/tensorflow_probability/python/math/psd_kernels/schur_complement.py
