Enable AutoCompositeTensor for `LinearOperator' bijectors.

PiperOrigin-RevId: 375531427

Author: emilyfertig

tensorflow_probability/python/bijectors/affine.py

@@ -213,7 +213,8 @@ def _as_tensor(x, name, dtype):
   return None if x is None else tf.convert_to_tensor(x, name=name, dtype=dtype)
 
 
-class Affine(bijector.Bijector):
+@bijector.auto_composite_tensor_bijector
+class Affine(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X; shift, scale) = scale @ X + shift`.
 
   Here `scale = c * I + diag(D1) + tril(L) + V @ diag(D2) @ V.T`.

tensorflow_probability/python/bijectors/affine_linear_operator.py

@@ -30,7 +30,8 @@
 ]
 
 
-class AffineLinearOperator(bijector.Bijector):
+@bijector.auto_composite_tensor_bijector
+class AffineLinearOperator(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X; shift, scale) = scale @ X + shift`.
 
   `shift` is a numeric `Tensor` and `scale` is a `LinearOperator`.

tensorflow_probability/python/bijectors/bijector_properties_test.py

@@ -196,11 +196,9 @@
 })
 COMPOSITE_TENSOR_ATOL = collections.defaultdict(lambda: 1e-6)
 
-# TODO(b/182603117): Enable AutoCT for meta-bijectors and LinearOperator.
+# TODO(b/182603117): Enable AutoCT for meta-bijectors.
 AUTO_COMPOSITE_TENSOR_IS_BROKEN = [
     'FillScaleTriL',
-    'ScaleMatvecDiag',
-    'ScaleMatvecTriL',
 ]
 
 

tensorflow_probability/python/bijectors/scale_matvec_diag.py

@@ -20,6 +20,7 @@
 
 import tensorflow.compat.v2 as tf
 
+from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import ldj_ratio
 from tensorflow_probability.python.bijectors import scale_matvec_linear_operator
 from tensorflow_probability.python.internal import dtype_util
@@ -32,8 +33,7 @@
 ]
 
 
-# TODO(b/182603117): Enable AutoCompositeTensor once LinearOperators are
-# converted to CompositeTensor.
+@bijector.auto_composite_tensor_bijector
 class ScaleMatvecDiag(scale_matvec_linear_operator.ScaleMatvecLinearOperator):
   """Compute `Y = g(X; scale) = scale @ X`.
 

tensorflow_probability/python/bijectors/scale_matvec_linear_operator.py

@@ -33,9 +33,7 @@
 ]
 
 
-# TODO(b/182603117): Enable AutoCompositeTensor once LinearOperators are
-# converted to CompositeTensor.
-class _ScaleMatvecLinearOperatorBase(bijector.Bijector):
+class _ScaleMatvecLinearOperatorBase(bijector.AutoCompositeTensorBijector):
   """Common base class for `ScaleMatvecLinearOperator{Block}`."""
 
   @property
@@ -68,6 +66,7 @@ def _parameter_control_dependencies(self, is_init):
     return []
 
 
+@bijector.auto_composite_tensor_bijector
 class ScaleMatvecLinearOperator(_ScaleMatvecLinearOperatorBase):
   """Compute `Y = g(X; scale) = scale @ X`.
 
@@ -144,6 +143,7 @@ def __init__(self,
           name=name)
 
 
+@bijector.auto_composite_tensor_bijector
 class ScaleMatvecLinearOperatorBlock(_ScaleMatvecLinearOperatorBase):
   """Compute `Y = g(X; scale) = scale @ X` for blockwise `X` and `scale`.
 

tensorflow_probability/python/bijectors/scale_matvec_tril.py

@@ -19,6 +19,7 @@
 from __future__ import print_function
 
 import tensorflow.compat.v2 as tf
+from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import fill_triangular as fill_triangular_bijector
 from tensorflow_probability.python.bijectors import scale_matvec_linear_operator
 from tensorflow_probability.python.internal import dtype_util
@@ -31,6 +32,7 @@
 ]
 
 
+@bijector.auto_composite_tensor_bijector
 class ScaleMatvecTriL(scale_matvec_linear_operator.ScaleMatvecLinearOperator):
   """Compute `Y = g(X; scale) = scale @ X`.
 

tensorflow_probability/python/experimental/vi/util/trainable_linear_operators.py

@@ -305,7 +305,8 @@ def build_linear_operator_zeros(
         is_self_adjoint=is_square, dtype=dtype)
 
 
-class _DefaultScaleDiagonal(bijector_lib.Bijector):
+@bijector_lib.auto_composite_tensor_bijector
+class _DefaultScaleDiagonal(bijector_lib.AutoCompositeTensorBijector):
   """Default bijector for constraining the diagonal of scale matrices."""
 
   def __init__(self):