Remove auto_composite_tensor decorators on bijector subclasses now that it's handled in a metaclass.

PiperOrigin-RevId: 377152250

Author: emilyfertig

tensorflow_probability/python/bijectors/BUILD

@@ -264,27 +264,26 @@ multi_substrate_py_library(
     name = "scale_matvec_diag",
     srcs = ["scale_matvec_diag.py"],
     deps = [
-        ":bijector",
         ":ldj_ratio",
+        ":scale_matvec_linear_operator",
         # tensorflow dep,
-        "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
+        "//tensorflow_probability/python/internal:tensor_util",
     ],
 )
 
 multi_substrate_py_library(
     name = "scale_matvec_tril",
     srcs = ["scale_matvec_tril.py"],
     deps = [
-        ":bijector",
         ":fill_triangular",
-        ":identity",
+        ":scale_matvec_linear_operator",
         # tensorflow dep,
-        "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:prefer_static",
+        "//tensorflow_probability/python/internal:tensor_util",
     ],
 )
 
@@ -301,11 +300,9 @@ multi_substrate_py_library(
     name = "blockwise",
     srcs = ["blockwise.py"],
     deps = [
-        ":bijector",
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:prefer_static",
         "//tensorflow_probability/python/internal:tensorshape_util",
     ],
@@ -414,12 +411,8 @@ multi_substrate_py_library(
     name = "exp",
     srcs = ["exp.py"],
     deps = [
-        ":bijector",
         ":invert",
         ":power_transform",
-        "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
-        "//tensorflow_probability/python/internal:prefer_static",
     ],
 )
 
@@ -430,7 +423,6 @@ multi_substrate_py_library(
         ":bijector",
         ":invert",
         # tensorflow dep,
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
     ],
 )
 
@@ -578,7 +570,6 @@ multi_substrate_py_library(
     srcs = ["invert.py"],
     deps = [
         ":bijector",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
     ],
 )
 

tensorflow_probability/python/bijectors/absolute_value.py

@@ -29,7 +29,6 @@
 ]
 
 
-@bijector.auto_composite_tensor_bijector
 class AbsoluteValue(bijector.AutoCompositeTensorBijector):
   """Computes `Y = g(X) = Abs(X)`, element-wise.
 

tensorflow_probability/python/bijectors/affine.py

@@ -213,7 +213,6 @@ def _as_tensor(x, name, dtype):
   return None if x is None else tf.convert_to_tensor(x, name=name, dtype=dtype)
 
 
-@bijector.auto_composite_tensor_bijector
 class Affine(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X; shift, scale) = scale @ X + shift`.
 

tensorflow_probability/python/bijectors/affine_linear_operator.py

@@ -30,7 +30,6 @@
 ]
 
 
-@bijector.auto_composite_tensor_bijector
 class AffineLinearOperator(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X; shift, scale) = scale @ X + shift`.
 

tensorflow_probability/python/bijectors/ascending.py

@@ -29,7 +29,6 @@
 ]
 
 
-@bijector.auto_composite_tensor_bijector
 class Ascending(bijector.AutoCompositeTensorBijector):
   """Maps unconstrained R^n to R^n in ascending order.
 

tensorflow_probability/python/bijectors/bijector.py

@@ -20,7 +20,6 @@
 
 import abc
 import contextlib
-import functools
 
 # Dependency imports
 import numpy as np
@@ -1778,13 +1777,6 @@ def _composite_tensor_shape_params(self):
       return ()
 
 
-auto_composite_tensor_bijector = functools.partial(
-    auto_composite_tensor.auto_composite_tensor,
-    omit_kwargs=('parameters',),
-    non_identifying_kwargs=('name',),
-    module_name='tfp.bijectors')
-
-
 class _AutoCompositeTensorBijectorMeta(_BijectorMeta):
   """Metaclass for `AutoCompositeTensorBijector`."""
 
@@ -1793,7 +1785,11 @@ def __new__(mcs, classname, baseclasses, attrs):  # pylint: disable=bad-mcs-clas
 
     cls = super(_AutoCompositeTensorBijectorMeta, mcs).__new__(  # pylint: disable=too-many-function-args
         mcs, classname, baseclasses, attrs)
-    return auto_composite_tensor_bijector(cls)
+    return auto_composite_tensor.auto_composite_tensor(
+        cls,
+        omit_kwargs=('parameters',),
+        non_identifying_kwargs=('name',),
+        module_name='tfp.bijectors')
 
 
 class AutoCompositeTensorBijector(

tensorflow_probability/python/bijectors/blockwise.py

@@ -28,7 +28,6 @@
 from tensorflow_probability.python.bijectors import joint_map
 from tensorflow_probability.python.bijectors import split
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import prefer_static as ps
 from tensorflow_probability.python.internal import tensorshape_util
 
@@ -270,8 +269,7 @@ def _validate_block_sizes(block_sizes, bijectors, validate_args):
   return block_sizes
 
 
-@bijector_lib.auto_composite_tensor_bijector
-class Blockwise(_Blockwise, auto_composite_tensor.AutoCompositeTensor):
+class Blockwise(_Blockwise, bijector_lib.AutoCompositeTensorBijector):
 
   def __new__(cls, *args, **kwargs):
     """Returns a `_Blockwise` if any of `bijectors` is not `CompositeTensor."""

tensorflow_probability/python/bijectors/categorical_to_discrete.py

@@ -36,7 +36,6 @@
 ]
 
 
-@bijector.auto_composite_tensor_bijector
 class CategoricalToDiscrete(bijector.AutoCompositeTensorBijector):
   """Bijector which computes `Y = g(X) = values[X]`.
 

tensorflow_probability/python/bijectors/cholesky_outer_product.py

@@ -35,7 +35,6 @@
 ]
 
 
-@bijector.auto_composite_tensor_bijector
 class CholeskyOuterProduct(bijector.AutoCompositeTensorBijector):
   """Compute `g(X) = X @ X.T`; X is lower-triangular, positive-diagonal matrix.
 

tensorflow_probability/python/bijectors/cholesky_to_inv_cholesky.py

@@ -32,7 +32,6 @@
 ]
 
 
-@bijector.auto_composite_tensor_bijector
 class CholeskyToInvCholesky(bijector.AutoCompositeTensorBijector):
   """Maps the Cholesky factor of `M` to the Cholesky factor of `M^{-1}`.