From 1b1beb793480cedfb9abf0db9039f570e870b085 Mon Sep 17 00:00:00 2001
From: emilyaf <emilyaf@google.com>
Date: Wed, 19 May 2021 14:34:56 -0700
Subject: [PATCH] Add an `auto_composite_tensor_bijector` decorator for
 bijectors that preserves the `name` attribute through flattening/unflattening
 and in serialization.

PiperOrigin-RevId: 374731367
---
 tensorflow_probability/python/bijectors/BUILD |  48 ------------------
 .../python/bijectors/absolute_value.py        |   4 +-
 .../python/bijectors/ascending.py             |   4 +-
 .../python/bijectors/bijector.py              |   8 +++
 .../python/bijectors/bijector_test.py         |   4 +-
 .../bijectors/categorical_to_discrete.py      |   3 +-
 .../bijectors/cholesky_outer_product.py       |   4 +-
 .../bijectors/cholesky_to_inv_cholesky.py     |   4 +-
 .../python/bijectors/correlation_cholesky.py  |   4 +-
 .../python/bijectors/cumsum.py                |   4 +-
 .../bijectors/discrete_cosine_transform.py    |   4 +-
 .../python/bijectors/exp.py                   |   7 ++-
 .../python/bijectors/expm1.py                 |   6 +--
 .../python/bijectors/fill_triangular.py       |   4 +-
 .../python/bijectors/frechet_cdf.py           |   4 +-
 .../python/bijectors/generalized_pareto.py    |   4 +-
 .../python/bijectors/gev_cdf.py               |   4 +-
 .../python/bijectors/gompertz_cdf.py          |   4 +-
 .../python/bijectors/gumbel_cdf.py            |   4 +-
 .../python/bijectors/identity.py              |   4 +-
 .../python/bijectors/inline.py                |   4 +-
 .../python/bijectors/invert.py                |   3 +-
 .../bijectors/iterated_sigmoid_centered.py    |   4 +-
 .../python/bijectors/kumaraswamy_cdf.py       |   4 +-
 .../python/bijectors/matrix_inverse_tril.py   |   4 +-
 .../python/bijectors/moyal_cdf.py             |   4 +-
 .../python/bijectors/normal_cdf.py            |   4 +-
 .../python/bijectors/ordered.py               |   4 +-
 .../python/bijectors/pad.py                   |   4 +-
 .../python/bijectors/permute.py               |   4 +-
 .../python/bijectors/power.py                 |   4 +-
 .../python/bijectors/power_transform.py       |   4 +-
 .../bijectors/rational_quadratic_spline.py    |   4 +-
 .../python/bijectors/rayleigh_cdf.py          |   4 +-
 .../python/bijectors/reciprocal.py            |   4 +-
 .../python/bijectors/reshape.py               |   4 +-
 .../python/bijectors/restructure.py           |   4 +-
 .../python/bijectors/scale.py                 |   4 +-
 .../python/bijectors/scale_matvec_lu.py       |   4 +-
 .../python/bijectors/shift.py                 |   4 +-
 .../python/bijectors/shifted_gompertz_cdf.py  |   4 +-
 .../python/bijectors/sigmoid.py               |   4 +-
 .../python/bijectors/sinh.py                  |   4 +-
 .../python/bijectors/sinh_arcsinh.py          |   4 +-
 .../python/bijectors/soft_clip.py             |   4 +-
 .../python/bijectors/softfloor.py             |   4 +-
 .../python/bijectors/softmax_centered.py      |   4 +-
 .../python/bijectors/softplus.py              |   4 +-
 .../python/bijectors/softsign.py              |   4 +-
 .../python/bijectors/split.py                 |   4 +-
 .../python/bijectors/square.py                |   4 +-
 .../python/bijectors/tanh.py                  |   4 +-
 .../python/bijectors/transform_diagonal.py    |   4 +-
 .../python/bijectors/transpose.py             |   4 +-
 .../python/bijectors/weibull_cdf.py           |   4 +-
 .../auto_composite_tensor/saved_model.pb      | Bin 17342 -> 17729 bytes
 .../python/util/deferred_tensor_test.py       |   4 +-
 57 files changed, 68 insertions(+), 207 deletions(-)

diff --git a/tensorflow_probability/python/bijectors/BUILD b/tensorflow_probability/python/bijectors/BUILD
index 0c6ae316c6..81dff46d00 100644
--- a/tensorflow_probability/python/bijectors/BUILD
+++ b/tensorflow_probability/python/bijectors/BUILD
@@ -167,7 +167,6 @@ multi_substrate_py_library(
         ":bijector",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:distribution_util",
     ],
 )
@@ -214,7 +213,6 @@ multi_substrate_py_library(
     deps = [
         ":bijector",
         # tensorflow dep,
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -227,7 +225,6 @@ multi_substrate_py_library(
     deps = [
         ":bijector",
         # tensorflow dep,
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -242,7 +239,6 @@ multi_substrate_py_library(
         ":softplus",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -359,7 +355,6 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:prefer_static",
@@ -375,7 +370,6 @@ multi_substrate_py_library(
         ":cholesky_outer_product",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:prefer_static",
     ],
@@ -388,7 +382,6 @@ multi_substrate_py_library(
         ":bijector",
         ":fill_triangular",
         # tensorflow dep,
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:prefer_static",
         "//tensorflow_probability/python/internal:tensorshape_util",
     ],
@@ -400,7 +393,6 @@ multi_substrate_py_library(
     deps = [
         ":bijector",
         # tensorflow dep,
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:prefer_static",
     ],
 )
@@ -412,7 +404,6 @@ multi_substrate_py_library(
         ":bijector",
         # tensorflow dep,
 #         "//tensorflow_probability/google:platform_google",  # DisableOnExport
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
     ],
 )
@@ -463,7 +454,6 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:tensorshape_util",
         "//tensorflow_probability/python/math:linalg",
@@ -478,7 +468,6 @@ multi_substrate_py_library(
         ":softplus",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -495,7 +484,6 @@ multi_substrate_py_library(
         ":sigmoid",
         ":softplus",
         # tensorflow dep,
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -531,7 +519,6 @@ multi_substrate_py_library(
         ":softplus",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -546,7 +533,6 @@ multi_substrate_py_library(
         ":softplus",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -561,7 +547,6 @@ multi_substrate_py_library(
         ":softplus",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -574,7 +559,6 @@ multi_substrate_py_library(
     deps = [
         ":bijector",
         # tensorflow dep,
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
     ],
 )
 
@@ -584,7 +568,6 @@ multi_substrate_py_library(
     deps = [
         ":bijector",
         # tensorflow dep,
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
     ],
 )
 
@@ -606,7 +589,6 @@ multi_substrate_py_library(
         ":softplus",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
@@ -660,7 +642,6 @@ multi_substrate_py_library(
         ":bijector",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
     ],
 )
@@ -674,7 +655,6 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -703,7 +683,6 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:special_math",
     ],
@@ -716,7 +695,6 @@ multi_substrate_py_library(
         ":bijector",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:distribution_util",
     ],
 )
@@ -728,7 +706,6 @@ multi_substrate_py_library(
         ":bijector",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:distribution_util",
     ],
 )
@@ -741,7 +718,6 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:prefer_static",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -757,7 +733,6 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
@@ -772,7 +747,6 @@ multi_substrate_py_library(
         ":bijector",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
     ],
@@ -785,7 +759,6 @@ multi_substrate_py_library(
         ":bijector",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
     ],
 )
 
@@ -798,7 +771,6 @@ multi_substrate_py_library(
         ":softplus",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -815,7 +787,6 @@ multi_substrate_py_library(
         ":softplus",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
     ],
@@ -840,7 +811,6 @@ multi_substrate_py_library(
         ":bijector",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:distribution_util",
     ],
 )
@@ -853,7 +823,6 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:nest_util",
         "//tensorflow_probability/python/internal:parameter_properties",
@@ -870,7 +839,6 @@ multi_substrate_py_library(
         ":bijector",
         ":invert",
         # tensorflow dep,
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:nest_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:prefer_static",
@@ -899,7 +867,6 @@ multi_substrate_py_library(
         ":bijector",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
     ],
@@ -911,7 +878,6 @@ multi_substrate_py_library(
     deps = [
         ":bijector",
         # tensorflow dep,
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/math",
     ],
 )
@@ -926,7 +892,6 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -941,7 +906,6 @@ multi_substrate_py_library(
         ":bijector",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -957,7 +921,6 @@ multi_substrate_py_library(
         ":softplus",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/math",
@@ -973,7 +936,6 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:prefer_static",
     ],
@@ -986,7 +948,6 @@ multi_substrate_py_library(
         ":bijector",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
@@ -1001,7 +962,6 @@ multi_substrate_py_library(
         ":bijector",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:dtype_util",
     ],
@@ -1015,7 +975,6 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:prefer_static",
         "//tensorflow_probability/python/internal:tensor_util",
@@ -1031,7 +990,6 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:distribution_util",
     ],
 )
@@ -1043,7 +1001,6 @@ multi_substrate_py_library(
         ":bijector",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:prefer_static",
     ],
@@ -1056,7 +1013,6 @@ multi_substrate_py_library(
         ":bijector",
         # numpy dep,
         # tensorflow dep,
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
     ],
 )
 
@@ -1066,7 +1022,6 @@ multi_substrate_py_library(
     deps = [
         ":bijector",
         # tensorflow dep,
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
     ],
 )
 
@@ -1078,7 +1033,6 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:distribution_util",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
@@ -1094,7 +1048,6 @@ multi_substrate_py_library(
         ":softplus",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
     ],
@@ -1142,7 +1095,6 @@ multi_substrate_py_test(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/distributions",
-        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:cache_util",
         "//tensorflow_probability/python/internal:parameter_properties",
         "//tensorflow_probability/python/internal:tensor_util",
diff --git a/tensorflow_probability/python/bijectors/absolute_value.py b/tensorflow_probability/python/bijectors/absolute_value.py
index 735398c4ae..3c84a85156 100644
--- a/tensorflow_probability/python/bijectors/absolute_value.py
+++ b/tensorflow_probability/python/bijectors/absolute_value.py
@@ -22,7 +22,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 
 __all__ = [
@@ -30,8 +29,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class AbsoluteValue(bijector.AutoCompositeTensorBijector):
   """Computes `Y = g(X) = Abs(X)`, element-wise.
 
diff --git a/tensorflow_probability/python/bijectors/ascending.py b/tensorflow_probability/python/bijectors/ascending.py
index 0f2223d466..be3594445c 100644
--- a/tensorflow_probability/python/bijectors/ascending.py
+++ b/tensorflow_probability/python/bijectors/ascending.py
@@ -22,7 +22,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 
 
 __all__ = [
@@ -30,8 +29,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Ascending(bijector.AutoCompositeTensorBijector):
   """Maps unconstrained R^n to R^n in ascending order.
 
diff --git a/tensorflow_probability/python/bijectors/bijector.py b/tensorflow_probability/python/bijectors/bijector.py
index b5df1d41d1..8cb9515488 100644
--- a/tensorflow_probability/python/bijectors/bijector.py
+++ b/tensorflow_probability/python/bijectors/bijector.py
@@ -20,6 +20,7 @@
 
 import abc
 import contextlib
+import functools
 
 # Dependency imports
 import numpy as np
@@ -1620,6 +1621,13 @@ class MyBijector(tfb.AutoCompositeTensorBijector):
   pass
 
 
+auto_composite_tensor_bijector = functools.partial(
+    auto_composite_tensor.auto_composite_tensor,
+    omit_kwargs=('parameters',),
+    non_identifying_kwargs=('name',),
+    module_name='tfp.bijectors')
+
+
 def check_valid_ndims(ndims, validate=True):
   """Ensures that `ndims` is a non-negative integer.
 
diff --git a/tensorflow_probability/python/bijectors/bijector_test.py b/tensorflow_probability/python/bijectors/bijector_test.py
index cb45c1fae1..82d4d34b39 100644
--- a/tensorflow_probability/python/bijectors/bijector_test.py
+++ b/tensorflow_probability/python/bijectors/bijector_test.py
@@ -25,7 +25,7 @@
 import tensorflow.compat.v1 as tf1
 import tensorflow.compat.v2 as tf
 from tensorflow_probability.python import bijectors as tfb
-from tensorflow_probability.python.internal import auto_composite_tensor
+from tensorflow_probability.python.bijectors import bijector as bijector_lib
 from tensorflow_probability.python.internal import cache_util
 from tensorflow_probability.python.internal import tensor_util
 from tensorflow_probability.python.internal import test_util
@@ -768,7 +768,7 @@ def testNestedCondition(self):
       mock_method.assert_called_once_with(mock.ANY, arg1=arg1, arg2=arg2)
 
 
-@auto_composite_tensor.auto_composite_tensor(omit_kwargs=('name',))
+@bijector_lib.auto_composite_tensor_bijector
 class CompositeForwardBijector(tfb.AutoCompositeTensorBijector):
 
   def __init__(self, scale=2., validate_args=False, name=None):
diff --git a/tensorflow_probability/python/bijectors/categorical_to_discrete.py b/tensorflow_probability/python/bijectors/categorical_to_discrete.py
index 91a4426de4..e762bf52bb 100644
--- a/tensorflow_probability/python/bijectors/categorical_to_discrete.py
+++ b/tensorflow_probability/python/bijectors/categorical_to_discrete.py
@@ -36,7 +36,8 @@
 ]
 
 
-class CategoricalToDiscrete(bijector.Bijector):
+@bijector.auto_composite_tensor_bijector
+class CategoricalToDiscrete(bijector.AutoCompositeTensorBijector):
   """Bijector which computes `Y = g(X) = values[X]`.
 
   Example Usage:
diff --git a/tensorflow_probability/python/bijectors/cholesky_outer_product.py b/tensorflow_probability/python/bijectors/cholesky_outer_product.py
index e826aacf28..9b1ebc7780 100644
--- a/tensorflow_probability/python/bijectors/cholesky_outer_product.py
+++ b/tensorflow_probability/python/bijectors/cholesky_outer_product.py
@@ -24,7 +24,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import distribution_util
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import prefer_static as ps
@@ -36,8 +35,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class CholeskyOuterProduct(bijector.AutoCompositeTensorBijector):
   """Compute `g(X) = X @ X.T`; X is lower-triangular, positive-diagonal matrix.
 
diff --git a/tensorflow_probability/python/bijectors/cholesky_to_inv_cholesky.py b/tensorflow_probability/python/bijectors/cholesky_to_inv_cholesky.py
index f1ec24ddac..c3cdebad11 100644
--- a/tensorflow_probability/python/bijectors/cholesky_to_inv_cholesky.py
+++ b/tensorflow_probability/python/bijectors/cholesky_to_inv_cholesky.py
@@ -23,7 +23,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors.cholesky_outer_product import CholeskyOuterProduct
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import prefer_static as ps
 
@@ -33,8 +32,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class CholeskyToInvCholesky(bijector.AutoCompositeTensorBijector):
   """Maps the Cholesky factor of `M` to the Cholesky factor of `M^{-1}`.
 
diff --git a/tensorflow_probability/python/bijectors/correlation_cholesky.py b/tensorflow_probability/python/bijectors/correlation_cholesky.py
index c83985a23d..1f5e425f8f 100644
--- a/tensorflow_probability/python/bijectors/correlation_cholesky.py
+++ b/tensorflow_probability/python/bijectors/correlation_cholesky.py
@@ -24,7 +24,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import fill_triangular
-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
 
@@ -33,8 +32,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class CorrelationCholesky(bijector.AutoCompositeTensorBijector):
   """Maps unconstrained reals to Cholesky-space correlation matrices.
 
diff --git a/tensorflow_probability/python/bijectors/cumsum.py b/tensorflow_probability/python/bijectors/cumsum.py
index 0023840500..1f89d661fc 100644
--- a/tensorflow_probability/python/bijectors/cumsum.py
+++ b/tensorflow_probability/python/bijectors/cumsum.py
@@ -20,7 +20,6 @@
 
 import tensorflow.compat.v2 as tf
 from tensorflow_probability.python.bijectors import bijector
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import prefer_static
 
 __all__ = [
@@ -28,8 +27,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Cumsum(bijector.AutoCompositeTensorBijector):
   """Computes the cumulative sum of a tensor along a specified axis.
 
diff --git a/tensorflow_probability/python/bijectors/discrete_cosine_transform.py b/tensorflow_probability/python/bijectors/discrete_cosine_transform.py
index d66e62349d..c9a96eb897 100644
--- a/tensorflow_probability/python/bijectors/discrete_cosine_transform.py
+++ b/tensorflow_probability/python/bijectors/discrete_cosine_transform.py
@@ -21,7 +21,6 @@
 import tensorflow.compat.v2 as tf
 
 from tensorflow_probability.python.bijectors import bijector
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 
 
@@ -30,8 +29,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class DiscreteCosineTransform(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X) = DCT(X)`, where DCT type is indicated by the `type` arg.
 
diff --git a/tensorflow_probability/python/bijectors/exp.py b/tensorflow_probability/python/bijectors/exp.py
index 0fe2565118..6bc0524fab 100644
--- a/tensorflow_probability/python/bijectors/exp.py
+++ b/tensorflow_probability/python/bijectors/exp.py
@@ -20,6 +20,7 @@
 
 import tensorflow.compat.v2 as tf
 
+from tensorflow_probability.python.bijectors import bijector as bijector_lib
 from tensorflow_probability.python.bijectors import invert
 from tensorflow_probability.python.bijectors import power_transform
 from tensorflow_probability.python.internal import auto_composite_tensor
@@ -31,8 +32,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector_lib.auto_composite_tensor_bijector
 class Exp(power_transform.PowerTransform):
   """Compute `Y = g(X) = exp(X)`.
 
@@ -76,8 +76,7 @@ def __init__(self,
 
 # TODO(b/182603117): Remove `AutoCompositeTensor` when `Invert` subclasses
 # `AutoCompositeTensor` and ensure `tf.saved_model` still works.
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector_lib.auto_composite_tensor_bijector
 class Log(invert.Invert,
           auto_composite_tensor.AutoCompositeTensor):
   """Compute `Y = log(X)`. This is `Invert(Exp())`."""
diff --git a/tensorflow_probability/python/bijectors/expm1.py b/tensorflow_probability/python/bijectors/expm1.py
index d51b644678..ea9d022c70 100644
--- a/tensorflow_probability/python/bijectors/expm1.py
+++ b/tensorflow_probability/python/bijectors/expm1.py
@@ -30,8 +30,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Expm1(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X) = exp(X) - 1`.
 
@@ -95,8 +94,7 @@ def _forward_log_det_jacobian(self, x):
 
 # TODO(b/182603117): Remove `AutoCompositeTensor` when `Invert` subclasses
 # `AutoCompositeTensor`.
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Log1p(invert.Invert, auto_composite_tensor.AutoCompositeTensor):
   """Compute `Y = log1p(X)`. This is `Invert(Expm1())`."""
 
diff --git a/tensorflow_probability/python/bijectors/fill_triangular.py b/tensorflow_probability/python/bijectors/fill_triangular.py
index ac3ba881dc..84f1643515 100644
--- a/tensorflow_probability/python/bijectors/fill_triangular.py
+++ b/tensorflow_probability/python/bijectors/fill_triangular.py
@@ -25,7 +25,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import tensorshape_util
 from tensorflow_probability.python.math.linalg import fill_triangular
 from tensorflow_probability.python.math.linalg import fill_triangular_inverse
@@ -36,8 +35,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class FillTriangular(bijector.AutoCompositeTensorBijector):
   """Transforms vectors to triangular.
 
diff --git a/tensorflow_probability/python/bijectors/frechet_cdf.py b/tensorflow_probability/python/bijectors/frechet_cdf.py
index 9f8297b527..24f4261f4b 100644
--- a/tensorflow_probability/python/bijectors/frechet_cdf.py
+++ b/tensorflow_probability/python/bijectors/frechet_cdf.py
@@ -23,7 +23,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import softplus as softplus_bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -34,8 +33,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class FrechetCDF(bijector.AutoCompositeTensorBijector):
   """The Frechet cumulative density function.
 
diff --git a/tensorflow_probability/python/bijectors/generalized_pareto.py b/tensorflow_probability/python/bijectors/generalized_pareto.py
index 6f4910589e..b5a75d04d3 100644
--- a/tensorflow_probability/python/bijectors/generalized_pareto.py
+++ b/tensorflow_probability/python/bijectors/generalized_pareto.py
@@ -24,7 +24,6 @@
 from tensorflow_probability.python.bijectors import shift as shift_bijector
 from tensorflow_probability.python.bijectors import sigmoid as sigmoid_bijector
 from tensorflow_probability.python.bijectors import softplus as softplus_bijector
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -35,8 +34,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector_lib.auto_composite_tensor_bijector
 class GeneralizedPareto(bijector_lib.AutoCompositeTensorBijector):
   """Bijector mapping R**n to non-negative reals.
 
diff --git a/tensorflow_probability/python/bijectors/gev_cdf.py b/tensorflow_probability/python/bijectors/gev_cdf.py
index 11cfbdd767..00ce4c4cc3 100644
--- a/tensorflow_probability/python/bijectors/gev_cdf.py
+++ b/tensorflow_probability/python/bijectors/gev_cdf.py
@@ -23,7 +23,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import softplus as softplus_bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -33,8 +32,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class GeneralizedExtremeValueCDF(bijector.AutoCompositeTensorBijector):
   """Compute the GeneralizedExtremeValue CDF.
 
diff --git a/tensorflow_probability/python/bijectors/gompertz_cdf.py b/tensorflow_probability/python/bijectors/gompertz_cdf.py
index 334de5d027..ebe493ee41 100644
--- a/tensorflow_probability/python/bijectors/gompertz_cdf.py
+++ b/tensorflow_probability/python/bijectors/gompertz_cdf.py
@@ -23,7 +23,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import softplus as softplus_bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -34,8 +33,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class GompertzCDF(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X) = 1 - exp(-c * (exp(rate * X) - 1)`, the Gompertz CDF.
 
diff --git a/tensorflow_probability/python/bijectors/gumbel_cdf.py b/tensorflow_probability/python/bijectors/gumbel_cdf.py
index aab68232e8..f1d215eef5 100644
--- a/tensorflow_probability/python/bijectors/gumbel_cdf.py
+++ b/tensorflow_probability/python/bijectors/gumbel_cdf.py
@@ -23,7 +23,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import softplus as softplus_bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -34,8 +33,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class GumbelCDF(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X) = exp(-exp(-(X - loc) / scale))`, the Gumbel CDF.
 
diff --git a/tensorflow_probability/python/bijectors/identity.py b/tensorflow_probability/python/bijectors/identity.py
index a7d350ce1a..add76e7121 100644
--- a/tensorflow_probability/python/bijectors/identity.py
+++ b/tensorflow_probability/python/bijectors/identity.py
@@ -21,7 +21,6 @@
 import tensorflow.compat.v2 as tf
 
 from tensorflow_probability.python.bijectors import bijector
-from tensorflow_probability.python.internal import auto_composite_tensor
 
 __all__ = [
     'Identity',
@@ -34,8 +33,7 @@ def __getitem__(self, _):
     return {}
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Identity(bijector.AutoCompositeTensorBijector):
   """Compute Y = g(X) = X.
 
diff --git a/tensorflow_probability/python/bijectors/inline.py b/tensorflow_probability/python/bijectors/inline.py
index b60a307ea1..9b5e8a5ece 100644
--- a/tensorflow_probability/python/bijectors/inline.py
+++ b/tensorflow_probability/python/bijectors/inline.py
@@ -21,7 +21,6 @@
 import tensorflow.compat.v2 as tf
 
 from tensorflow_probability.python.bijectors import bijector
-from tensorflow_probability.python.internal import auto_composite_tensor
 
 
 __all__ = [
@@ -29,8 +28,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Inline(bijector.AutoCompositeTensorBijector):
   """Bijector constructed from custom callables.
 
diff --git a/tensorflow_probability/python/bijectors/invert.py b/tensorflow_probability/python/bijectors/invert.py
index feab4295ab..02a1541365 100644
--- a/tensorflow_probability/python/bijectors/invert.py
+++ b/tensorflow_probability/python/bijectors/invert.py
@@ -138,8 +138,7 @@ def forward_event_ndims(self, event_ndims, **kwargs):
     return self.bijector.inverse_event_ndims(event_ndims, **kwargs)
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name', 'parameters'), module_name='tfp.bijectors')
+@bijector_lib.auto_composite_tensor_bijector
 class Invert(_Invert, auto_composite_tensor.AutoCompositeTensor):
 
   def __new__(cls, *args, **kwargs):
diff --git a/tensorflow_probability/python/bijectors/iterated_sigmoid_centered.py b/tensorflow_probability/python/bijectors/iterated_sigmoid_centered.py
index 30ef831ced..a8ac21a49c 100644
--- a/tensorflow_probability/python/bijectors/iterated_sigmoid_centered.py
+++ b/tensorflow_probability/python/bijectors/iterated_sigmoid_centered.py
@@ -22,7 +22,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import prefer_static as ps
 
@@ -32,8 +31,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class IteratedSigmoidCentered(bijector.AutoCompositeTensorBijector):
   """Bijector which applies a Stick Breaking procedure.
 
diff --git a/tensorflow_probability/python/bijectors/kumaraswamy_cdf.py b/tensorflow_probability/python/bijectors/kumaraswamy_cdf.py
index 79a77ebd7e..f961fadd06 100644
--- a/tensorflow_probability/python/bijectors/kumaraswamy_cdf.py
+++ b/tensorflow_probability/python/bijectors/kumaraswamy_cdf.py
@@ -23,7 +23,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import softplus as softplus_bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import distribution_util
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
@@ -35,8 +34,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class KumaraswamyCDF(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X) = (1 - X**a)**b, X in [0, 1]`.
 
diff --git a/tensorflow_probability/python/bijectors/matrix_inverse_tril.py b/tensorflow_probability/python/bijectors/matrix_inverse_tril.py
index 88f0b8785c..15a4f856ff 100644
--- a/tensorflow_probability/python/bijectors/matrix_inverse_tril.py
+++ b/tensorflow_probability/python/bijectors/matrix_inverse_tril.py
@@ -22,7 +22,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 
 
@@ -31,8 +30,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class MatrixInverseTriL(bijector.AutoCompositeTensorBijector):
   """Computes `g(L) = inv(L)`, where `L` is a lower-triangular matrix.
 
diff --git a/tensorflow_probability/python/bijectors/moyal_cdf.py b/tensorflow_probability/python/bijectors/moyal_cdf.py
index 265f71ae1b..89aeb97719 100644
--- a/tensorflow_probability/python/bijectors/moyal_cdf.py
+++ b/tensorflow_probability/python/bijectors/moyal_cdf.py
@@ -27,7 +27,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import softplus as softplus_bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -38,8 +37,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class MoyalCDF(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X) = erfc(exp(- 1/2 * (X - loc) / scale) / sqrt(2))`.
 
diff --git a/tensorflow_probability/python/bijectors/normal_cdf.py b/tensorflow_probability/python/bijectors/normal_cdf.py
index abcfea99d0..3ca3369a5e 100644
--- a/tensorflow_probability/python/bijectors/normal_cdf.py
+++ b/tensorflow_probability/python/bijectors/normal_cdf.py
@@ -25,7 +25,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import special_math
 
@@ -34,8 +33,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class NormalCDF(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X) = NormalCDF(x)`.
 
diff --git a/tensorflow_probability/python/bijectors/ordered.py b/tensorflow_probability/python/bijectors/ordered.py
index 4bad67c8f3..046a0c2c53 100644
--- a/tensorflow_probability/python/bijectors/ordered.py
+++ b/tensorflow_probability/python/bijectors/ordered.py
@@ -22,7 +22,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow.python.util import deprecation  # pylint: disable=g-direct-tensorflow-import
 
 
@@ -31,8 +30,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Ordered(bijector.AutoCompositeTensorBijector):
   """Maps a vector of increasing elements to an unconstrained vector.
 
diff --git a/tensorflow_probability/python/bijectors/pad.py b/tensorflow_probability/python/bijectors/pad.py
index 813089fa45..d8c423c1d5 100644
--- a/tensorflow_probability/python/bijectors/pad.py
+++ b/tensorflow_probability/python/bijectors/pad.py
@@ -23,7 +23,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import prefer_static as ps
 from tensorflow_probability.python.internal import tensor_util
@@ -35,8 +34,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Pad(bijector.AutoCompositeTensorBijector):
   """Pads a value to the `event_shape` of a `Tensor`.
 
diff --git a/tensorflow_probability/python/bijectors/permute.py b/tensorflow_probability/python/bijectors/permute.py
index accf30b591..c2412f92eb 100644
--- a/tensorflow_probability/python/bijectors/permute.py
+++ b/tensorflow_probability/python/bijectors/permute.py
@@ -25,7 +25,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -37,8 +36,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Permute(bijector.AutoCompositeTensorBijector):
   """Permutes the rightmost dimension of a `Tensor`.
 
diff --git a/tensorflow_probability/python/bijectors/power.py b/tensorflow_probability/python/bijectors/power.py
index e595e9b409..b9b628cb83 100644
--- a/tensorflow_probability/python/bijectors/power.py
+++ b/tensorflow_probability/python/bijectors/power.py
@@ -24,7 +24,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import prefer_static as ps
 from tensorflow_probability.python.internal import tensor_util
@@ -38,8 +37,7 @@ def _is_odd_integer(x):
   return ps.equal(x, ps.round(x)) & ps.not_equal(2. * ps.floor(x / 2.), x)
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Power(bijector.AutoCompositeTensorBijector):
   """Compute `g(X) = X ** power`; where X is a non-negative real number.
 
diff --git a/tensorflow_probability/python/bijectors/power_transform.py b/tensorflow_probability/python/bijectors/power_transform.py
index cd653d9fe3..d96d945772 100644
--- a/tensorflow_probability/python/bijectors/power_transform.py
+++ b/tensorflow_probability/python/bijectors/power_transform.py
@@ -22,7 +22,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 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
 
 
@@ -31,8 +30,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class PowerTransform(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X) = (1 + X * c)**(1 / c), X >= -1 / c`.
 
diff --git a/tensorflow_probability/python/bijectors/rational_quadratic_spline.py b/tensorflow_probability/python/bijectors/rational_quadratic_spline.py
index eb3454a533..eab94194c8 100644
--- a/tensorflow_probability/python/bijectors/rational_quadratic_spline.py
+++ b/tensorflow_probability/python/bijectors/rational_quadratic_spline.py
@@ -26,7 +26,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import softplus as softplus_bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -59,8 +58,7 @@ def _knot_positions(bin_sizes, range_min):
     'SplineShared', 'out_of_bounds,x_k,y_k,d_k,d_kp1,h_k,w_k,s_k')
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class RationalQuadraticSpline(bijector.AutoCompositeTensorBijector):
   """A piecewise rational quadratic spline, as developed in [1].
 
diff --git a/tensorflow_probability/python/bijectors/rayleigh_cdf.py b/tensorflow_probability/python/bijectors/rayleigh_cdf.py
index f36fcfade9..22ca6de0e8 100644
--- a/tensorflow_probability/python/bijectors/rayleigh_cdf.py
+++ b/tensorflow_probability/python/bijectors/rayleigh_cdf.py
@@ -21,7 +21,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import softplus as softplus_bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -31,8 +30,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class RayleighCDF(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X) = 1 - exp( -(X/scale)**2 / 2 ), X >= 0`.
 
diff --git a/tensorflow_probability/python/bijectors/reciprocal.py b/tensorflow_probability/python/bijectors/reciprocal.py
index e5fca69ae2..40eda26d9d 100644
--- a/tensorflow_probability/python/bijectors/reciprocal.py
+++ b/tensorflow_probability/python/bijectors/reciprocal.py
@@ -22,14 +22,12 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 
 __all__ = ['Reciprocal']
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Reciprocal(bijector.AutoCompositeTensorBijector):
   """A `Bijector` that computes the reciprocal `b(x) = 1. / x` entrywise.
 
diff --git a/tensorflow_probability/python/bijectors/reshape.py b/tensorflow_probability/python/bijectors/reshape.py
index 0a56d9193f..362b8f4d9b 100644
--- a/tensorflow_probability/python/bijectors/reshape.py
+++ b/tensorflow_probability/python/bijectors/reshape.py
@@ -25,7 +25,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import prefer_static as ps
@@ -38,8 +37,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Reshape(bijector.AutoCompositeTensorBijector):
   """Reshapes the `event_shape` of a `Tensor`.
 
diff --git a/tensorflow_probability/python/bijectors/restructure.py b/tensorflow_probability/python/bijectors/restructure.py
index 98a8b3fd1f..cb6b2b13c9 100644
--- a/tensorflow_probability/python/bijectors/restructure.py
+++ b/tensorflow_probability/python/bijectors/restructure.py
@@ -24,7 +24,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import invert
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import nest_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow.python.util import nest  # pylint: disable=g-direct-tensorflow-import
@@ -45,8 +44,7 @@ def unique_token_set(source_structure):
   return flat_token_set
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Restructure(bijector.AutoCompositeTensorBijector):
   """Converts between nested structures of Tensors.
 
diff --git a/tensorflow_probability/python/bijectors/scale.py b/tensorflow_probability/python/bijectors/scale.py
index 0402f15e64..ac6cb662de 100644
--- a/tensorflow_probability/python/bijectors/scale.py
+++ b/tensorflow_probability/python/bijectors/scale.py
@@ -22,7 +22,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -33,8 +32,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Scale(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X; scale) = scale * X`.
 
diff --git a/tensorflow_probability/python/bijectors/scale_matvec_lu.py b/tensorflow_probability/python/bijectors/scale_matvec_lu.py
index 0f93a0ddfc..f9b6b19158 100644
--- a/tensorflow_probability/python/bijectors/scale_matvec_lu.py
+++ b/tensorflow_probability/python/bijectors/scale_matvec_lu.py
@@ -22,7 +22,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import prefer_static as ps
 from tensorflow_probability.python.internal import tensor_util
@@ -38,8 +37,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class ScaleMatvecLU(bijector.AutoCompositeTensorBijector):
   """Matrix-vector multiply using LU decomposition.
 
diff --git a/tensorflow_probability/python/bijectors/shift.py b/tensorflow_probability/python/bijectors/shift.py
index 38962dd5df..c11fc1228d 100644
--- a/tensorflow_probability/python/bijectors/shift.py
+++ b/tensorflow_probability/python/bijectors/shift.py
@@ -21,7 +21,6 @@
 import tensorflow.compat.v2 as tf
 
 from tensorflow_probability.python.bijectors import bijector
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -32,8 +31,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Shift(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X; shift) = X + shift`.
 
diff --git a/tensorflow_probability/python/bijectors/shifted_gompertz_cdf.py b/tensorflow_probability/python/bijectors/shifted_gompertz_cdf.py
index c18a25e40b..90f30f83e9 100644
--- a/tensorflow_probability/python/bijectors/shifted_gompertz_cdf.py
+++ b/tensorflow_probability/python/bijectors/shifted_gompertz_cdf.py
@@ -25,7 +25,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import softplus as softplus_bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -36,8 +35,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class ShiftedGompertzCDF(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X) = (1 - exp(-rate * X)) * exp(-c * exp(-rate * X))`.
 
diff --git a/tensorflow_probability/python/bijectors/sigmoid.py b/tensorflow_probability/python/bijectors/sigmoid.py
index 48b3a9af87..839ad3b0ab 100644
--- a/tensorflow_probability/python/bijectors/sigmoid.py
+++ b/tensorflow_probability/python/bijectors/sigmoid.py
@@ -21,7 +21,6 @@
 import tensorflow.compat.v2 as tf
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -65,8 +64,7 @@ def grad_fn(dy):
     return y, grad_fn
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Sigmoid(bijector.AutoCompositeTensorBijector):
   """Bijector that computes the logistic sigmoid function.
 
diff --git a/tensorflow_probability/python/bijectors/sinh.py b/tensorflow_probability/python/bijectors/sinh.py
index c9f1c14a65..9e5b1606ab 100644
--- a/tensorflow_probability/python/bijectors/sinh.py
+++ b/tensorflow_probability/python/bijectors/sinh.py
@@ -21,7 +21,6 @@
 import tensorflow.compat.v2 as tf
 from tensorflow_probability.python import math as tfp_math
 from tensorflow_probability.python.bijectors import bijector
-from tensorflow_probability.python.internal import auto_composite_tensor
 
 
 __all__ = [
@@ -29,8 +28,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Sinh(bijector.AutoCompositeTensorBijector):
   """Bijector that computes `Y = sinh(X)`.
 
diff --git a/tensorflow_probability/python/bijectors/sinh_arcsinh.py b/tensorflow_probability/python/bijectors/sinh_arcsinh.py
index 8262963477..2aeadf0718 100644
--- a/tensorflow_probability/python/bijectors/sinh_arcsinh.py
+++ b/tensorflow_probability/python/bijectors/sinh_arcsinh.py
@@ -24,7 +24,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import softplus as softplus_bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -34,8 +33,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class SinhArcsinh(bijector.AutoCompositeTensorBijector):
   """`Y = g(X) = Sinh( (Arcsinh(X) + skewness) * tailweight ) * multiplier`.
 
diff --git a/tensorflow_probability/python/bijectors/soft_clip.py b/tensorflow_probability/python/bijectors/soft_clip.py
index 0a4d45b6dd..563d9074ec 100644
--- a/tensorflow_probability/python/bijectors/soft_clip.py
+++ b/tensorflow_probability/python/bijectors/soft_clip.py
@@ -31,7 +31,6 @@
 from tensorflow_probability.python.bijectors import softplus
 
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -42,8 +41,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class SoftClip(bijector.AutoCompositeTensorBijector):
   """Bijector that approximates clipping as a continuous, differentiable map.
 
diff --git a/tensorflow_probability/python/bijectors/softfloor.py b/tensorflow_probability/python/bijectors/softfloor.py
index a8e213bd0e..6df7a679f1 100644
--- a/tensorflow_probability/python/bijectors/softfloor.py
+++ b/tensorflow_probability/python/bijectors/softfloor.py
@@ -26,7 +26,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import softplus as softplus_bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -37,8 +36,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Softfloor(bijector.AutoCompositeTensorBijector):
   """Compute a differentiable approximation to `tf.math.floor`.
 
diff --git a/tensorflow_probability/python/bijectors/softmax_centered.py b/tensorflow_probability/python/bijectors/softmax_centered.py
index 5d4d0b63fc..b5e69c9812 100644
--- a/tensorflow_probability/python/bijectors/softmax_centered.py
+++ b/tensorflow_probability/python/bijectors/softmax_centered.py
@@ -25,7 +25,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import pad as pad_lib
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import prefer_static as ps
 
@@ -35,8 +34,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class SoftmaxCentered(bijector.AutoCompositeTensorBijector):
   """Bijector which computes `Y = g(X) = exp([X 0]) / sum(exp([X 0]))`.
 
diff --git a/tensorflow_probability/python/bijectors/softplus.py b/tensorflow_probability/python/bijectors/softplus.py
index 9b7de51502..92abae943c 100644
--- a/tensorflow_probability/python/bijectors/softplus.py
+++ b/tensorflow_probability/python/bijectors/softplus.py
@@ -21,7 +21,6 @@
 import tensorflow.compat.v2 as tf
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import distribution_util
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
@@ -59,8 +58,7 @@ def grad_fn(dy):
     return y, grad_fn
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Softplus(bijector.AutoCompositeTensorBijector):
   """Bijector which computes `Y = g(X) = Log[1 + exp(X)]`.
 
diff --git a/tensorflow_probability/python/bijectors/softsign.py b/tensorflow_probability/python/bijectors/softsign.py
index d7cdb076e5..f388dca1a8 100644
--- a/tensorflow_probability/python/bijectors/softsign.py
+++ b/tensorflow_probability/python/bijectors/softsign.py
@@ -22,7 +22,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 
 
@@ -31,8 +30,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Softsign(bijector.AutoCompositeTensorBijector):
   """Bijector which computes `Y = g(X) = X / (1 + |X|)`.
 
diff --git a/tensorflow_probability/python/bijectors/split.py b/tensorflow_probability/python/bijectors/split.py
index 87f98f3719..1bed48136b 100644
--- a/tensorflow_probability/python/bijectors/split.py
+++ b/tensorflow_probability/python/bijectors/split.py
@@ -25,7 +25,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import prefer_static
 from tensorflow_probability.python.internal import tensor_util
@@ -36,8 +35,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Split(bijector.AutoCompositeTensorBijector):
   """Split a `Tensor` event along an axis into a list of `Tensor`s.
 
diff --git a/tensorflow_probability/python/bijectors/square.py b/tensorflow_probability/python/bijectors/square.py
index ed7f91f6b0..5ec5168da6 100644
--- a/tensorflow_probability/python/bijectors/square.py
+++ b/tensorflow_probability/python/bijectors/square.py
@@ -25,7 +25,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 
 
 __all__ = [
@@ -33,8 +32,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Square(bijector.AutoCompositeTensorBijector):
   """Compute `g(X) = X^2`; X is a positive real number.
 
diff --git a/tensorflow_probability/python/bijectors/tanh.py b/tensorflow_probability/python/bijectors/tanh.py
index 57b8309916..2483348760 100644
--- a/tensorflow_probability/python/bijectors/tanh.py
+++ b/tensorflow_probability/python/bijectors/tanh.py
@@ -21,7 +21,6 @@
 import numpy as np
 import tensorflow.compat.v2 as tf
 from tensorflow_probability.python.bijectors import bijector
-from tensorflow_probability.python.internal import auto_composite_tensor
 
 
 __all__ = [
@@ -29,8 +28,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Tanh(bijector.AutoCompositeTensorBijector):
   """Bijector that computes `Y = tanh(X)`, therefore `Y in (-1, 1)`.
 
diff --git a/tensorflow_probability/python/bijectors/transform_diagonal.py b/tensorflow_probability/python/bijectors/transform_diagonal.py
index 7df7ad6d2c..fc527ea475 100644
--- a/tensorflow_probability/python/bijectors/transform_diagonal.py
+++ b/tensorflow_probability/python/bijectors/transform_diagonal.py
@@ -20,15 +20,13 @@
 
 import tensorflow.compat.v2 as tf
 from tensorflow_probability.python.bijectors import bijector
-from tensorflow_probability.python.internal import auto_composite_tensor
 
 __all__ = [
     'TransformDiagonal',
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name', 'parameters'), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class TransformDiagonal(bijector.AutoCompositeTensorBijector):
   """Applies a Bijector to the diagonal of a matrix.
 
diff --git a/tensorflow_probability/python/bijectors/transpose.py b/tensorflow_probability/python/bijectors/transpose.py
index 67d11931f1..b6dfdd91a0 100644
--- a/tensorflow_probability/python/bijectors/transpose.py
+++ b/tensorflow_probability/python/bijectors/transpose.py
@@ -25,7 +25,6 @@
 
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import distribution_util
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
@@ -38,8 +37,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class Transpose(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X) = transpose_rightmost_dims(X, rightmost_perm)`.
 
diff --git a/tensorflow_probability/python/bijectors/weibull_cdf.py b/tensorflow_probability/python/bijectors/weibull_cdf.py
index c20ec003aa..8734aa6aee 100644
--- a/tensorflow_probability/python/bijectors/weibull_cdf.py
+++ b/tensorflow_probability/python/bijectors/weibull_cdf.py
@@ -23,7 +23,6 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import softplus as softplus_bijector
 from tensorflow_probability.python.internal import assert_util
-from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import dtype_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import tensor_util
@@ -34,8 +33,7 @@
 ]
 
 
-@auto_composite_tensor.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='tfp.bijectors')
+@bijector.auto_composite_tensor_bijector
 class WeibullCDF(bijector.AutoCompositeTensorBijector):
   """Compute `Y = g(X) = 1 - exp( -( X / scale) ** concentration), X >= 0`.
 
diff --git a/tensorflow_probability/python/internal/testdata/auto_composite_tensor/saved_model.pb b/tensorflow_probability/python/internal/testdata/auto_composite_tensor/saved_model.pb
index 3e61c7c1d41dfe1ac705d3869a2755c233540342..c9d6e0640e34f221b4c1aabc7f2c64cc5341b6c1 100644
GIT binary patch
delta 1439
zcmchWNoZ3+7{~j~<R!1)GZrUJlSo~X?yUrkiD|VPg7hFBJk*o$e5SEB#H42NQXzuk
z5=&jalNLeir9@oXr-xFk6czNMtzL>k!GqvM#DaM7;7eM>;?kUYoEc{3_y5g*rq021
z79<(OWOhvQTG-$*Jz#0z+#hJxHM{`a-|W{vE1N*}w*<wghI6#Bs=Bber%D5}-{0bM
zA<d$To>%w7QQA>cSLMM<RSWf~N7AuSsxKZ>L-AND5s#|zbgDm{609C*w!%RZju?HY
zxy${&X8lU-OOxfTf)@cd(~MG1@n9)8?$if@w@eVG4XvJ{&+Tq~V8&r5W;uOrs-%N<
z8x1vCEPh{$=*In|HLa)nwKm$k#iGwN1xwjH(B;5(Y7Mj%b;XCez53I@6_Do2k$_$)
zl?6-adM!hU-nokg)rCg^>6-&@0}WL>t%gh?Aslvs@FLAyoqD%*s=(c7M=NQD6|cQZ
zGIaU8bkpY5N9}ik4FMf+>??}gkgiWSwt(c3<kM!tw(*K549k+?#EHIGui6tnqNSrL
zafi{HwwlZ*KzZ@S)UDz@G|mA&0C|nb!U#VL$PSs&mE#pyj#4tDMZ-eLGFg}elbI8+
zGsoMo3`?W&UUhMqVF}vN6w{7{g*BoSCRLxhP2Gkzw4TtSeLY$#tZIqgq+tC}gX>24
zT3i#sel20_Nre+h<RiQUl@TtZIZ`-@;_V^Rpn_d=Qm&+h_5srZ^L?S?9WE)abCxc3
z3}+r0#AAT>;3~`I6Pz)|r&dg`GslB}CMa7<V2LcI{ww=(bjQEh-89)*m;X-mvGWGI
z&htIU<iUS10kD4B0DTO&y8QKq{N(}u$KSr&oEbO7-ph+UZ}6U3QEYpT@A^}$ZAq*W
zDg9NABIFuu<^L##hh1aju5wF;U$H>ncQ<5SK#o5Hoc_Jq4-NWNDt2R`*+t9!334Wg
A`2YX_

delta 910
zcmb`FJxo(^6vesk-^Y8s72E!yrIto%V<jb~dGtd*Y^ewdE*cwSbRc;=@Im-0$k(EA
zGN~f5i8rFA8l`L%g8CSYF(`v38iOPbZf-6{TpS!6^tF={i|i-ocg{U`V-D7{Aj-g>
zXP=5;N%%Zs#HA2p%8{s{;dNk2ORKSMI{-~hN`(JZajO=W0QaftP-Z}%)HD6MI;M^2
zYTwYX8dc&kE+X`KfQON8Mm&ZXS(hXw8f$c66=}8_`qJ7hD$!^pOs8AI#`(4tfMJpn
zC#j(9Hx?2Hp%T|>kbpLAzA}8i;iC*G>QMqlLGgoStrBM`iT-lLqg9$2F&?9bd*`8x
z4%-__KD3Y+px^ct6CCB|@C>Evy!6QFGrl@<7V%dt&M`{Y_$cS7rYm*jeEck?%d%c)
zH935e@w%?r3NL8G*+!9Iv+3$KmV$3UTD9Q~z<Sy;yX0I}W}d~!yB00bK}m0*q@#P@
zCt^U9e~9!d`OHlM+Nbo%*`ds!dO;u1W=5y^AE0b|u>K{WJU%l$F*B_yd=(m(0k49*
zhpPqFgF>%`Hrsu!P>~%%J6cD_2i2*5ZB*w1R*<99?X(q`jFrQdj5ektE3gV$?C{gY
z)N#u<D|XXm*+-vJmjWMz$U9-D#U9hdvAib-1!gw82DgO&4RNBh4Snx;(D{Bx@Up-G
vym=FR2XGhp0XHkMiokAE@iV>+rFRp}O8hOI>Iwz#Ly_GDy#4QbyTbYlBX#_a

diff --git a/tensorflow_probability/python/util/deferred_tensor_test.py b/tensorflow_probability/python/util/deferred_tensor_test.py
index a2b57d8b94..e1e824de1d 100644
--- a/tensorflow_probability/python/util/deferred_tensor_test.py
+++ b/tensorflow_probability/python/util/deferred_tensor_test.py
@@ -622,8 +622,8 @@ def testRepr(self):
         "_TransformedVariableSpec(input_spec=TensorSpec(shape=(4, 2), "
         "dtype=tf.float32, name=None), "
         "transform_or_spec=Sigmoid_ACTTypeSpec(3, {}, {'low': None, 'high': "
-        "None, 'validate_args': True}, ('name',), (), (), {}), dtype=<dtype: "
-        "'float64'>, name=None)")
+        "None, 'validate_args': True, 'name': 'sigmoid'}, ('parameters',), (), "
+        "('name',), {}), dtype=<dtype: 'float64'>, name=None)")
     self.assertEqual(repr(spec), expected)