Repository: tensorflow/addons

Issue:
Generated documentation of type aliases are included in unrelated modules
For example, the documentation of `TensorLike` is located in `tfa.image.color_ops`:

https://www.tensorflow.org/addons/api_docs/python/tfa/image/color_ops/TensorLike

Presumably, the system selects the module where the symbol was first found. This is not a big issue but it produces strange and unexpected documentation for type annotations, see for example:

https://www.tensorflow.org/addons/api_docs/python/tfa/seq2seq/LuongAttention

I see 2 possible solutions:

1. Hide non public symbols in the generated documentation
1. Properly export type aliases in a public module

Discussion:
WindQAQ: @MarkDaoust, Hi Mark, can you help review this when time allows? Thank you!
MarkDaoust: > explicit_package_contents_filter

Good idea. 

The main thing to remember with this setting is that the now the imports in your`__init__.py` are the definition of your public API. With this is's best not to put any **code** in the init files other than imports.

If that's what you want -> great. If ever something isn't showing up it's probably just missing an import.

If you've been test-running this **could you post a `git diff --stat` of the output directory**? this would make it clearwhat affect this has on the generated pages.
WindQAQ: > The main thing to remember with this setting is that the now the imports in your__init__.py are the definition of your public API. With this is's best not to put any code in the init files other than imports.

That's actually what I want. /cc @seanpmorgan @bhack for discussion.

Here is the diff file.

<details><summary>diff</summary>
<p>

```diff
 .../{addons_api => new_addons_api}/_redirects.yaml |  444 +-
 .../{addons_api => new_addons_api}/_toc.yaml       |  384 +-
 .../{addons_api => new_addons_api}/index.md        |  342 +-
 .../tfa/_api_cache.json                            | 6081 +-------------------
 .../tfa/activations/gelu.md                        |    2 +-
 .../tfa/activations/hardshrink.md                  |    6 +-
 .../tfa/activations/lisht.md                       |    2 +-
 .../tfa/activations/mish.md                        |    2 +-
 .../tfa/activations/rrelu.md                       |    6 +-
 .../tfa/activations/snake.md                       |    4 +-
 .../tfa/activations/softshrink.md                  |    6 +-
 .../tfa/activations/sparsemax.md                   |    2 +-
 .../tfa/activations/tanhshrink.md                  |    2 +-
 .../tfa/callbacks/AverageModelCheckpoint.md        |    8 -
 .../tfa/callbacks/TQDMProgressBar.md               |    8 -
 .../tfa/callbacks/TimeStopping.md                  |    8 -
 .../average_model_checkpoint.md => dev/null}       |   28 -
 .../tfa/callbacks/time_stopping.md => dev/null}    |   30 -
 .../callbacks/tqdm_progress_bar.md => dev/null}    |   30 -
 .../tfa/callbacks.md                               |    8 -
 .../tfa/image/adjust_hsv_in_yiq.md                 |   16 +-
 .../tfa/image/blend.md                             |   14 +-
 .../image/color_ops/equalize_image.md => dev/null} |   34 -
 .../color_ops/sharpness_image.md => dev/null}      |   34 -
 .../addons_api/tfa/image/color_ops.md => dev/null} |   43 -
 .../tfa/image/compose_ops.md => dev/null}          |   36 -
 .../tfa/image/connected_components.md              |    2 +-
 .../tfa/image/cutout.md                            |   16 +-
 .../tfa/image/cutout_ops.md => dev/null}           |   38 -
 .../tfa/image/dense_image_warp.md                  |    4 +-
 .../tfa/image/distance_transform.md => dev/null}   |   34 -
 .../tfa/image/distort_image_ops.md => dev/null}    |   38 -
 .../tfa/image/equalize.md                          |   10 +-
 .../tfa/image/euclidean_dist_transform.md          |   10 +-
 .../addons_api/tfa/image/filters.md => dev/null}   |   38 -
 .../tfa/image/gaussian_filter2d.md                 |   14 +-
 .../tfa/image/interpolate_bilinear.md              |    4 +-
 .../tfa/image/interpolate_spline.md                |    8 +-
 .../tfa/image/mean_filter2d.md                     |   12 +-
 .../tfa/image/median_filter2d.md                   |   12 +-
 .../tfa/image/random_cutout.md                     |   16 +-
 .../tfa/image/random_hsv_in_yiq.md                 |   20 +-
 .../tfa/image/resampler.md                         |   12 +-
 .../tfa/image/resampler_ops.md => dev/null}        |   30 -
 .../tfa/image/rotate.md                            |   12 +-
 .../tfa/image/sharpness.md                         |   12 +-
 .../tfa/image/shear_x.md                           |   14 +-
 .../tfa/image/shear_y.md                           |   14 +-
 .../tfa/image/sparse_image_warp.md                 |    8 +-
 .../tfa/image/transform.md                         |   12 +-
 .../null}                                          |   84 -
 .../compose_transforms.md => dev/null}             |   77 -
 .../flat_transforms_to_matrices.md => dev/null}    |   94 -
 .../matrices_to_flat_transforms.md => dev/null}    |   95 -
 .../tfa/image/transform_ops.md => dev/null}        |   48 -
 .../tfa/image/translate.md                         |   12 +-
 .../null}                                          |   76 -
 .../tfa/image/translate_ops.md => dev/null}        |   38 -
 .../tfa/image/translate_xy.md                      |   16 +-
 .../tfa/image/utils/from_4D_image.md => dev/null}  |   71 -
 .../tfa/image/utils/get_ndims.md => dev/null}      |   31 -
 .../tfa/image/utils/to_4D_image.md => dev/null}    |   64 -
 .../tfa/image/utils/unwrap.md => dev/null}         |   81 -
 .../tfa/image/utils/wrap.md => dev/null}           |   33 -
 .../addons_api/tfa/image/utils.md => dev/null}     |   38 -
 .../{addons_api => new_addons_api}/tfa/image.md    |   22 -
 .../tfa/layers/AdaptiveAveragePooling1D.md         |   10 -
 .../tfa/layers/AdaptiveAveragePooling2D.md         |   10 -
 .../tfa/layers/AdaptiveAveragePooling3D.md         |   10 -
 .../tfa/layers/AdaptiveMaxPooling1D.md             |   10 -
 .../tfa/layers/AdaptiveMaxPooling2D.md             |   10 -
 .../tfa/layers/AdaptiveMaxPooling3D.md             |   10 -
 .../tfa/layers/CorrelationCost.md                  |    8 -
 .../tfa/layers/ESN.md                              |   24 +-
 .../tfa/layers/FilterResponseNormalization.md      |   22 +-
 .../tfa/layers/GELU.md                             |    8 -
 .../tfa/layers/GroupNormalization.md               |   20 +-
 .../tfa/layers/InstanceNormalization.md            |    8 -
 .../tfa/layers/Maxout.md                           |    8 -
 .../tfa/layers/MultiHeadAttention.md               |    8 -
 .../tfa/layers/PoincareNormalize.md                |    8 -
 .../tfa/layers/PolynomialCrossing.md               |   16 +-
 .../tfa/layers/Snake.md                            |   10 +-
 .../tfa/layers/Sparsemax.md                        |    8 -
 .../tfa/layers/SpatialPyramidPooling2D.md          |    8 -
 .../tfa/layers/SpectralNormalization.md            |    8 -
 .../tfa/layers/TLU.md                              |   20 +-
 .../tfa/layers/WeightNormalization.md              |    8 -
 .../AdaptivePooling1D.md => dev/null}              | 1043 ----
 .../AdaptivePooling2D.md => dev/null}              | 1043 ----
 .../AdaptivePooling3D.md => dev/null}              | 1043 ----
 .../tfa/layers/adaptive_pooling.md => dev/null}    |   46 -
 .../addons_api/tfa/layers/esn.md => dev/null}      |   40 -
 .../addons_api/tfa/layers/gelu.md => dev/null}     |   30 -
 .../addons_api/tfa/layers/maxout.md => dev/null}   |   30 -
 .../tfa/layers/multihead_attention.md => dev/null} |   28 -
 .../tfa/layers/normalizations.md => dev/null}      |   32 -
 .../tfa/layers/optical_flow.md => dev/null}        |   30 -
 .../addons_api/tfa/layers/poincare.md => dev/null} |   30 -
 .../tfa/layers/polynomial.md => dev/null}          |   30 -
 .../addons_api/tfa/layers/snake.md => dev/null}    |   30 -
 .../tfa/layers/sparsemax.md => dev/null}           |   28 -
 .../layers/spatial_pyramid_pooling.md => dev/null} |   30 -
 .../layers/spectral_normalization.md => dev/null}  |   28 -
 .../addons_api/tfa/layers/tlu.md => dev/null}      |   30 -
 .../addons_api/tfa/layers/wrappers.md => dev/null} |   28 -
 .../{addons_api => new_addons_api}/tfa/layers.md   |   32 -
 .../tfa/losses/ContrastiveLoss.md                  |   10 +-
 .../tfa/losses/LiftedStructLoss.md                 |   10 +-
 .../tfa/losses/NpairsLoss.md                       |    8 -
 .../tfa/losses/NpairsMultilabelLoss.md             |    8 -
 .../tfa/losses/PinballLoss.md                      |   10 +-
 .../tfa/losses/SigmoidFocalCrossEntropy.md         |   12 +-
 .../tfa/losses/TripletHardLoss.md                  |   10 +-
 .../tfa/losses/TripletSemiHardLoss.md              |   10 +-
 .../tfa/losses/WeightedKappaLoss.md                |    8 -
 .../tfa/losses/contrastive.md => dev/null}         |   40 -
 .../tfa/losses/contrastive_loss.md                 |   14 +-
 .../tfa/losses/focal_loss.md => dev/null}          |   40 -
 .../tfa/losses/giou_loss.md                        |    4 +-
 .../tfa/losses/kappa_loss.md => dev/null}          |   34 -
 .../addons_api/tfa/losses/lifted.md => dev/null}   |   40 -
 .../tfa/losses/lifted_struct_loss.md               |   14 +-
 .../angular_distance.md => dev/null}               |   69 -
 .../pairwise_distance.md => dev/null}              |   77 -
 .../tfa/losses/metric_learning.md => dev/null}     |   36 -
 .../addons_api/tfa/losses/npairs.md => dev/null}   |   42 -
 .../tfa/losses/npairs_loss.md                      |   12 +-
 .../tfa/losses/npairs_multilabel_loss.md           |   12 +-
 .../tfa/losses/pinball_loss.md                     |   14 +-
 .../tfa/losses/quantiles.md => dev/null}           |   40 -
 .../tfa/losses/sigmoid_focal_crossentropy.md       |   16 +-
 .../tfa/losses/sparsemax_loss.md                   |    6 +-
 .../addons_api/tfa/losses/triplet.md => dev/null}  |   44 -
 .../tfa/losses/triplet_hard_loss.md                |   14 +-
 .../tfa/losses/triplet_semihard_loss.md            |   14 +-
 .../{addons_api => new_addons_api}/tfa/losses.md   |   18 -
 .../tfa/metrics/CohenKappa.md                      |   12 +-
 .../tfa/metrics/F1Score.md                         |   12 +-
 .../tfa/metrics/FBetaScore.md                      |   14 +-
 .../tfa/metrics/GeometricMean.md                   |   10 +-
 .../tfa/metrics/HammingLoss.md                     |   10 +-
 .../tfa/metrics/MatthewsCorrelationCoefficient.md  |   12 +-
 .../tfa/metrics/MeanMetricWrapper.md               |   10 +-
 .../tfa/metrics/MultiLabelConfusionMatrix.md       |   12 +-
 .../tfa/metrics/RSquare.md                         |   10 +-
 .../tfa/metrics/cohens_kappa.md => dev/null}       |   36 -
 .../tfa/metrics/f_scores.md => dev/null}           |   38 -
 .../tfa/metrics/geometric_mean.md => dev/null}     |   34 -
 .../hamming/hamming_loss_fn.md => dev/null}        |  128 -
 .../addons_api/tfa/metrics/hamming.md => dev/null} |   44 -
 .../tfa/metrics/hamming_distance.md                |   12 +-
 .../null}                                          |   36 -
 .../multilabel_confusion_matrix.md => dev/null}    |   36 -
 .../tfa/metrics/r_square.md => dev/null}           |   38 -
 .../sample_weight_shape_match.md => dev/null}      |   31 -
 .../addons_api/tfa/metrics/utils.md => dev/null}   |   38 -
 .../{addons_api => new_addons_api}/tfa/metrics.md  |   18 -
 .../tfa/optimizers/AdamW.md                        |   12 +-
 .../tfa/optimizers/AveragedOptimizerWrapper.md     |   10 +-
 .../tfa/optimizers/ConditionalGradient.md          |   10 +-
 .../tfa/optimizers/CyclicalLearningRate.md         |   10 +-
 .../optimizers/ExponentialCyclicalLearningRate.md  |   12 +-
 .../tfa/optimizers/LAMB.md                         |   16 +-
 .../tfa/optimizers/LazyAdam.md                     |   14 +-
 .../tfa/optimizers/Lookahead.md                    |   12 +-
 .../tfa/optimizers/MovingAverage.md                |   12 +-
 .../tfa/optimizers/NovoGrad.md                     |   16 +-
 .../tfa/optimizers/ProximalAdagrad.md              |    8 -
 .../tfa/optimizers/RectifiedAdam.md                |   20 +-
 .../tfa/optimizers/SGDW.md                         |   10 -
 .../tfa/optimizers/SWA.md                          |   10 +-
 .../optimizers/Triangular2CyclicalLearningRate.md  |   10 +-
 .../optimizers/TriangularCyclicalLearningRate.md   |   10 +-
 .../tfa/optimizers/Yogi.md                         |   20 +-
 .../tfa/optimizers/average_wrapper.md => dev/null} |   28 -
 .../conditional_gradient.md => dev/null}           |   34 -
 .../cyclical_learning_rate.md => dev/null}         |   40 -
 .../extend_with_decoupled_weight_decay.md          |    8 -
 .../addons_api/tfa/optimizers/lamb.md => dev/null} |   36 -
 .../tfa/optimizers/lazy_adam.md => dev/null}       |   38 -
 .../tfa/optimizers/lookahead.md => dev/null}       |   28 -
 .../tfa/optimizers/moving_average.md => dev/null}  |   28 -
 .../tfa/optimizers/novograd.md => dev/null}        |   34 -
 .../optimizers/proximal_adagrad.md => dev/null}    |   34 -
 .../tfa/optimizers/rectified_adam.md => dev/null}  |   34 -
 .../stochastic_weight_averaging.md => dev/null}    |   37 -
 .../DecoupledWeightDecayExtension.md => dev/null}  |  345 --
 .../weight_decay_optimizers.md => dev/null}        |   42 -
 .../addons_api/tfa/optimizers/yogi.md => dev/null} |   40 -
 .../tfa/optimizers.md                              |   28 -
 .../tfa/options/warn_fallback.md => dev/null}      |   31 -
 .../windqaq/addons_api/tfa/options.md => dev/null} |   34 -
 .../register_custom_kernels.md => dev/null}        |   29 -
 .../register_keras_objects.md => dev/null}         |   29 -
 .../addons_api/tfa/register.md => dev/null}        |   32 -
 .../tfa}/register_all.md                           |   14 +-
 .../tfa/rnn/ESNCell.md                             |   16 +-
 .../tfa/rnn/LayerNormLSTMCell.md                   |   42 +-
 .../tfa/rnn/LayerNormSimpleRNNCell.md              |   42 +-
 .../tfa/rnn/NASCell.md                             |   18 +-
 .../tfa/rnn/PeepholeLSTMCell.md                    |    8 -
 .../addons_api/tfa/rnn/esn_cell.md => dev/null}    |   36 -
 .../tfa/rnn/layer_norm_lstm_cell.md => dev/null}   |   44 -
 .../rnn/layer_norm_simple_rnn_cell.md => dev/null} |   44 -
 .../addons_api/tfa/rnn/nas_cell.md => dev/null}    |   38 -
 .../tfa/rnn/peephole_lstm_cell.md => dev/null}     |   30 -
 .../{addons_api => new_addons_api}/tfa/rnn.md      |   12 -
 .../tfa/seq2seq/AttentionMechanism.md              |    8 -
 .../tfa/seq2seq/AttentionWrapper.md                |    8 -
 .../tfa/seq2seq/AttentionWrapperState.md           |    8 -
 .../tfa/seq2seq/BahdanauAttention.md               |   14 +-
 .../tfa/seq2seq/BahdanauMonotonicAttention.md      |   18 +-
 .../tfa/seq2seq/BaseDecoder.md                     |    8 -
 .../tfa/seq2seq/BasicDecoder.md                    |    8 -
 .../tfa/seq2seq/BasicDecoderOutput.md              |    8 -
 .../tfa/seq2seq/BeamSearchDecoder.md               |   14 +-
 .../tfa/seq2seq/BeamSearchDecoderOutput.md         |    8 -
 .../tfa/seq2seq/BeamSearchDecoderState.md          |    8 -
 .../tfa/seq2seq/CustomSampler.md                   |   12 +-
 .../tfa/seq2seq/Decoder.md                         |    8 -
 .../tfa/seq2seq/FinalBeamSearchDecoderOutput.md    |    8 -
 .../tfa/seq2seq/GreedyEmbeddingSampler.md          |    8 -
 .../tfa/seq2seq/InferenceSampler.md                |   12 +-
 .../tfa/seq2seq/LuongAttention.md                  |   12 +-
 .../tfa/seq2seq/LuongMonotonicAttention.md         |   16 +-
 .../tfa/seq2seq/SampleEmbeddingSampler.md          |    8 -
 .../tfa/seq2seq/Sampler.md                         |    8 -
 .../seq2seq/ScheduledEmbeddingTrainingSampler.md   |   10 +-
 .../tfa/seq2seq/ScheduledOutputTrainingSampler.md  |   10 +-
 .../tfa/seq2seq/SequenceLoss.md                    |    8 -
 .../tfa/seq2seq/TrainingSampler.md                 |    8 -
 .../tfa/seq2seq/attention_wrapper.md => dev/null}  |   62 -
 .../tfa/seq2seq/basic_decoder.md => dev/null}      |   32 -
 .../BeamSearchDecoderMixin.md => dev/null}         |  318 -
 .../null}                                          |   66 -
 .../get_attention_probs.md => dev/null}            |   92 -
 .../seq2seq/beam_search_decoder.md => dev/null}    |   58 -
 .../addons_api/tfa/seq2seq/decoder.md => dev/null} |   40 -
 .../tfa/seq2seq/dynamic_decode.md                  |    8 -
 .../tfa/seq2seq/gather_tree.md                     |   16 +-
 .../tfa/seq2seq/gather_tree_from_array.md          |   14 +-
 .../tfa/seq2seq/hardmax.md                         |   10 +-
 .../addons_api/tfa/seq2seq/loss.md => dev/null}    |   38 -
 .../tfa/seq2seq/monotonic_attention.md             |   12 +-
 .../tfa/seq2seq/safe_cumprod.md                    |   10 +-
 .../sampler/bernoulli_sample.md => dev/null}       |   33 -
 .../sampler/categorical_sample.md => dev/null}     |   33 -
 .../addons_api/tfa/seq2seq/sampler.md => dev/null} |   56 -
 .../tfa/seq2seq/sequence_loss.md                   |   14 +-
 .../tfa/seq2seq/tile_batch.md                      |   10 +-
 .../{addons_api => new_addons_api}/tfa/seq2seq.md  |   14 -
 .../crf/CrfDecodeForwardRnnCell.md => dev/null}    | 1053 ----
 .../addons_api/tfa/text/crf.md => dev/null}        |   56 -
 .../tfa/text/crf_binary_score.md                   |   14 +-
 .../tfa/text/crf_decode.md                         |   14 +-
 .../tfa/text/crf_decode_backward.md                |   12 +-
 .../tfa/text/crf_decode_forward.md                 |   16 +-
 .../tfa/text/crf_forward.md                        |   16 +-
 .../tfa/text/crf_log_likelihood.md                 |   14 +-
 .../tfa/text/crf_log_norm.md                       |   14 +-
 .../tfa/text/crf_multitag_sequence_score.md        |   16 +-
 .../tfa/text/crf_sequence_score.md                 |   16 +-
 .../tfa/text/crf_unary_score.md                    |   14 +-
 .../tfa/text/parse_time.md                         |    8 -
 .../tfa/text/parse_time_op.md => dev/null}         |   34 -
 .../tfa/text/skip_gram_ops.md => dev/null}         |   40 -
 .../tfa/text/skip_gram_sample.md                   |   18 +-
 .../tfa/text/skip_gram_sample_with_text_vocab.md   |   22 +-
 .../tfa/text/viterbi_decode.md                     |   12 +-
 .../{addons_api => new_addons_api}/tfa/text.md     |    8 -
 .../tfa/types}/AcceptableDTypes.md                 |    4 +-
 .../esn => new_addons_api/tfa/types}/Activation.md |    4 +-
 .../tfa/types/Constraint.md}                       |    6 +-
 .../tfa/types}/FloatTensorLike.md                  |    4 +-
 .../tfa/types}/Number.md                           |    4 +-
 .../windqaq/new_addons_api/tfa/types/Optimizer.md} |   32 +
 .../tfa/types}/TensorLike.md                       |    4 +-
 .../windqaq/new_addons_api/tfa/types.md}           |   46 +
 .../windqaq/{addons_api => new_addons_api}/tfa.md  |    8 +-
 280 files changed, 485 insertions(+), 17753 deletions(-)

```

</p>
</details>

MarkDaoust: SGTM.

 Sorry that diff's not as clear as I was hoping... if you diff the `git diff` the generated _toc.yaml  that might show more clearly what's changed.
WindQAQ: > SGTM.
> 
> Sorry that diff's not as clear as I was hoping... if you diff the `git diff` the generated _toc.yaml that might show more clearly what's changed.

Here it is.

<details><summary>Diff TOC</summary>
<p>

```diff
diff --git a/home/windqaq/addons_api/_toc.yaml b/home/windqaq/new_addons_api/_toc.yaml
index 7c9c58a..45ee1f7 100644
--- a/home/windqaq/addons_api/_toc.yaml
+++ b/home/windqaq/new_addons_api/_toc.yaml
@@ -3,6 +3,8 @@ toc:
   section:
   - title: Overview
     path: /addons/api_docs/python/tfa
+  - title: register_all
+    path: /addons/api_docs/python/tfa/register_all
 - title: tfa.activations
   section:
   - title: Overview
@@ -35,18 +37,6 @@ toc:
     path: /addons/api_docs/python/tfa/callbacks/TimeStopping
   - title: TQDMProgressBar
     path: /addons/api_docs/python/tfa/callbacks/TQDMProgressBar
-  - title: average_model_checkpoint
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/callbacks/average_model_checkpoint
-  - title: time_stopping
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/callbacks/time_stopping
-  - title: tqdm_progress_bar
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/callbacks/tqdm_progress_bar
 - title: tfa.image
   section:
   - title: Overview
@@ -97,76 +87,6 @@ toc:
     path: /addons/api_docs/python/tfa/image/translate
   - title: translate_xy
     path: /addons/api_docs/python/tfa/image/translate_xy
-  - title: color_ops
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/image/color_ops
-    - title: equalize_image
-      path: /addons/api_docs/python/tfa/image/color_ops/equalize_image
-    - title: Number
-      path: /addons/api_docs/python/tfa/image/color_ops/Number
-    - title: sharpness_image
-      path: /addons/api_docs/python/tfa/image/color_ops/sharpness_image
-    - title: TensorLike
-      path: /addons/api_docs/python/tfa/image/color_ops/TensorLike
-  - title: compose_ops
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/image/compose_ops
-  - title: cutout_ops
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/image/cutout_ops
-  - title: distance_transform
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/image/distance_transform
-  - title: distort_image_ops
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/image/distort_image_ops
-  - title: filters
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/image/filters
-    - title: FloatTensorLike
-      path: /addons/api_docs/python/tfa/image/filters/FloatTensorLike
-  - title: resampler_ops
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/image/resampler_ops
-  - title: transform_ops
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/image/transform_ops
-    - title: angles_to_projective_transforms
-      path: /addons/api_docs/python/tfa/image/transform_ops/angles_to_projective_transforms
-    - title: compose_transforms
-      path: /addons/api_docs/python/tfa/image/transform_ops/compose_transforms
-    - title: flat_transforms_to_matrices
-      path: /addons/api_docs/python/tfa/image/transform_ops/flat_transforms_to_matrices
-    - title: matrices_to_flat_transforms
-      path: /addons/api_docs/python/tfa/image/transform_ops/matrices_to_flat_transforms
-  - title: translate_ops
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/image/translate_ops
-    - title: translations_to_projective_transforms
-      path: /addons/api_docs/python/tfa/image/translate_ops/translations_to_projective_transforms
-  - title: utils
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/image/utils
-    - title: from_4D_image
-      path: /addons/api_docs/python/tfa/image/utils/from_4D_image
-    - title: get_ndims
-      path: /addons/api_docs/python/tfa/image/utils/get_ndims
-    - title: to_4D_image
-      path: /addons/api_docs/python/tfa/image/utils/to_4D_image
-    - title: unwrap
-      path: /addons/api_docs/python/tfa/image/utils/unwrap
-    - title: wrap
-      path: /addons/api_docs/python/tfa/image/utils/wrap
 - title: tfa.layers
   section:
   - title: Overview
@@ -215,76 +135,6 @@ toc:
     path: /addons/api_docs/python/tfa/layers/TLU
   - title: WeightNormalization
     path: /addons/api_docs/python/tfa/layers/WeightNormalization
-  - title: adaptive_pooling
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/adaptive_pooling
-    - title: AdaptivePooling1D
-      path: /addons/api_docs/python/tfa/layers/adaptive_pooling/AdaptivePooling1D
-    - title: AdaptivePooling2D
-      path: /addons/api_docs/python/tfa/layers/adaptive_pooling/AdaptivePooling2D
-    - title: AdaptivePooling3D
-      path: /addons/api_docs/python/tfa/layers/adaptive_pooling/AdaptivePooling3D
-  - title: esn
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/esn
-    - title: Activation
-      path: /addons/api_docs/python/tfa/layers/esn/Activation
-    - title: Initializer
-      path: /addons/api_docs/python/tfa/layers/esn/Initializer
-  - title: gelu
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/gelu
-  - title: maxout
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/maxout
-  - title: multihead_attention
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/multihead_attention
-  - title: normalizations
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/normalizations
-  - title: optical_flow
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/optical_flow
-  - title: poincare
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/poincare
-  - title: polynomial
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/polynomial
-  - title: snake
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/snake
-  - title: sparsemax
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/sparsemax
-  - title: spatial_pyramid_pooling
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/spatial_pyramid_pooling
-  - title: spectral_normalization
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/spectral_normalization
-  - title: tlu
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/tlu
-  - title: wrappers
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/layers/wrappers
 - title: tfa.losses
   section:
   - title: Overview
@@ -331,42 +181,6 @@ toc:
     path: /addons/api_docs/python/tfa/losses/triplet_semihard_loss
   - title: WeightedKappaLoss
     path: /addons/api_docs/python/tfa/losses/WeightedKappaLoss
-  - title: contrastive
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/losses/contrastive
-  - title: focal_loss
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/losses/focal_loss
-  - title: kappa_loss
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/losses/kappa_loss
-  - title: lifted
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/losses/lifted
-  - title: metric_learning
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/losses/metric_learning
-    - title: angular_distance
-      path: /addons/api_docs/python/tfa/losses/metric_learning/angular_distance
-    - title: pairwise_distance
-      path: /addons/api_docs/python/tfa/losses/metric_learning/pairwise_distance
-  - title: npairs
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/losses/npairs
-  - title: quantiles
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/losses/quantiles
-  - title: triplet
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/losses/triplet
 - title: tfa.metrics
   section:
   - title: Overview
@@ -391,44 +205,6 @@ toc:
     path: /addons/api_docs/python/tfa/metrics/MultiLabelConfusionMatrix
   - title: RSquare
     path: /addons/api_docs/python/tfa/metrics/RSquare
-  - title: cohens_kappa
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/metrics/cohens_kappa
-    - title: AcceptableDTypes
-      path: /addons/api_docs/python/tfa/metrics/cohens_kappa/AcceptableDTypes
-  - title: f_scores
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/metrics/f_scores
-  - title: geometric_mean
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/metrics/geometric_mean
-  - title: hamming
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/metrics/hamming
-    - title: hamming_loss_fn
-      path: /addons/api_docs/python/tfa/metrics/hamming/hamming_loss_fn
-  - title: matthews_correlation_coefficient
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/metrics/matthews_correlation_coefficient
-  - title: multilabel_confusion_matrix
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/metrics/multilabel_confusion_matrix
-  - title: r_square
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/metrics/r_square
-  - title: utils
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/metrics/utils
-    - title: sample_weight_shape_match
-      path: /addons/api_docs/python/tfa/metrics/utils/sample_weight_shape_match
 - title: tfa.optimizers
   section:
   - title: Overview
@@ -469,76 +245,6 @@ toc:
     path: /addons/api_docs/python/tfa/optimizers/TriangularCyclicalLearningRate
   - title: Yogi
     path: /addons/api_docs/python/tfa/optimizers/Yogi
-  - title: average_wrapper
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/optimizers/average_wrapper
-  - title: conditional_gradient
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/optimizers/conditional_gradient
-  - title: cyclical_learning_rate
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/optimizers/cyclical_learning_rate
-  - title: lamb
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/optimizers/lamb
-  - title: lazy_adam
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/optimizers/lazy_adam
-  - title: lookahead
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/optimizers/lookahead
-  - title: moving_average
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/optimizers/moving_average
-  - title: novograd
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/optimizers/novograd
-  - title: proximal_adagrad
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/optimizers/proximal_adagrad
-  - title: rectified_adam
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/optimizers/rectified_adam
-  - title: stochastic_weight_averaging
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/optimizers/stochastic_weight_averaging
-  - title: weight_decay_optimizers
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/optimizers/weight_decay_optimizers
-    - title: DecoupledWeightDecayExtension
-      path: /addons/api_docs/python/tfa/optimizers/weight_decay_optimizers/DecoupledWeightDecayExtension
-  - title: yogi
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/optimizers/yogi
-- title: tfa.options
-  section:
-  - title: Overview
-    path: /addons/api_docs/python/tfa/options
-  - title: warn_fallback
-    path: /addons/api_docs/python/tfa/options/warn_fallback
-- title: tfa.register
-  section:
-  - title: Overview
-    path: /addons/api_docs/python/tfa/register
-  - title: register_all
-    path: /addons/api_docs/python/tfa/register/register_all
-  - title: register_custom_kernels
-    path: /addons/api_docs/python/tfa/register/register_custom_kernels
-  - title: register_keras_objects
-    path: /addons/api_docs/python/tfa/register/register_keras_objects
 - title: tfa.rnn
   section:
   - title: Overview
@@ -553,26 +259,6 @@ toc:
     path: /addons/api_docs/python/tfa/rnn/NASCell
   - title: PeepholeLSTMCell
     path: /addons/api_docs/python/tfa/rnn/PeepholeLSTMCell
-  - title: esn_cell
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/rnn/esn_cell
-  - title: layer_norm_lstm_cell
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/rnn/layer_norm_lstm_cell
-  - title: layer_norm_simple_rnn_cell
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/rnn/layer_norm_simple_rnn_cell
-  - title: nas_cell
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/rnn/nas_cell
-  - title: peephole_lstm_cell
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/rnn/peephole_lstm_cell
 - title: tfa.seq2seq
   section:
   - title: Overview
@@ -641,40 +327,6 @@ toc:
     path: /addons/api_docs/python/tfa/seq2seq/tile_batch
   - title: TrainingSampler
     path: /addons/api_docs/python/tfa/seq2seq/TrainingSampler
-  - title: attention_wrapper
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/seq2seq/attention_wrapper
-  - title: basic_decoder
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/seq2seq/basic_decoder
-  - title: beam_search_decoder
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/seq2seq/beam_search_decoder
-    - title: attention_probs_from_attn_state
-      path: /addons/api_docs/python/tfa/seq2seq/beam_search_decoder/attention_probs_from_attn_state
-    - title: BeamSearchDecoderMixin
-      path: /addons/api_docs/python/tfa/seq2seq/beam_search_decoder/BeamSearchDecoderMixin
-    - title: get_attention_probs
-      path: /addons/api_docs/python/tfa/seq2seq/beam_search_decoder/get_attention_probs
-  - title: decoder
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/seq2seq/decoder
-  - title: loss
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/seq2seq/loss
-  - title: sampler
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/seq2seq/sampler
-    - title: bernoulli_sample
-      path: /addons/api_docs/python/tfa/seq2seq/sampler/bernoulli_sample
-    - title: categorical_sample
-      path: /addons/api_docs/python/tfa/seq2seq/sampler/categorical_sample
 - title: tfa.text
   section:
   - title: Overview
@@ -707,17 +359,21 @@ toc:
     path: /addons/api_docs/python/tfa/text/skip_gram_sample_with_text_vocab
   - title: viterbi_decode
     path: /addons/api_docs/python/tfa/text/viterbi_decode
-  - title: crf
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/text/crf
-    - title: CrfDecodeForwardRnnCell
-      path: /addons/api_docs/python/tfa/text/crf/CrfDecodeForwardRnnCell
-  - title: parse_time_op
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/text/parse_time_op
-  - title: skip_gram_ops
-    section:
-    - title: Overview
-      path: /addons/api_docs/python/tfa/text/skip_gram_ops
+- title: tfa.types
+  section:
+  - title: Overview
+    path: /addons/api_docs/python/tfa/types
+  - title: AcceptableDTypes
+    path: /addons/api_docs/python/tfa/types/AcceptableDTypes
+  - title: Activation
+    path: /addons/api_docs/python/tfa/types/Activation
+  - title: Constraint
+    path: /addons/api_docs/python/tfa/types/Constraint
+  - title: FloatTensorLike
+    path: /addons/api_docs/python/tfa/types/FloatTensorLike
+  - title: Number
+    path: /addons/api_docs/python/tfa/types/Number
+  - title: Optimizer
+    path: /addons/api_docs/python/tfa/types/Optimizer
+  - title: TensorLike
+    path: /addons/api_docs/python/tfa/types/TensorLike
```

</p>
</details>
MarkDaoust: Cool thanks, 

So this looks mostly like a good change. But there are a few objects that will be dropped from the docs, many of these look like they are meant to be public, so if we are making this change you will need to add the missing imports to get **the ones that are meant to be public** to show up in the generated docs. 

```
tfa/image/color_ops
image/filters/FloatTensorLike
tfa/image/transform_ops
tfa/image/translate_ops/translations_to_projective_transforms
tfa/image/utils
tfa/losses/metric_learning
tfa/metrics/hamming/hamming_loss_fn
tfa/metrics/utils/sample_weight_shape_match
tfa/optimizers/weight_decay_optimizers/DecoupledWeightDecayExtension
tfa/options/warn_fallback
tfa/register
tfa/seq2seq/beam_search_decoder
tfa/seq2seq/sampler
tfa/text/crf
```

We could use some feedback from the addons team here to indicate which of these are meant for public consumption.
WindQAQ: > Cool thanks,
> 
> So this looks mostly like a good change. But there are a few objects that will be dropped from the docs, many of these look like they are meant to be public, so if we are making this change you will need to add the missing imports to get **the ones that are meant to be public** to show up in the generated docs.
> 
> ```
> tfa/image/color_ops
> image/filters/FloatTensorLike
> tfa/image/transform_ops
> tfa/image/translate_ops/translations_to_projective_transforms
> tfa/image/utils
> tfa/losses/metric_learning
> tfa/metrics/hamming/hamming_loss_fn
> tfa/metrics/utils/sample_weight_shape_match
> tfa/optimizers/weight_decay_optimizers/DecoupledWeightDecayExtension
> tfa/options/warn_fallback
> tfa/register
> tfa/seq2seq/beam_search_decoder
> tfa/seq2seq/sampler
> tfa/text/crf
> ```
> 
> We could use some feedback from the addons team here to indicate which of these are meant for public consumption.

Yep, some modules should be public. We'll clean it up and discuss it. Thank you!
WindQAQ: @tensorflow/sig-addons-maintainers I have listed all modules to be discussed down below. Any comment on this? Thank you!

- tfa/image/color_ops
    - **Drop**: `equalize` and `sharpness` already exported.
- image/filters/FloatTensorLike
    - **Drop**: this is just the alias of `tfa.utils.types.*`
- tfa/image/transform_ops
    - **Expose**
        - `compose_transforms`: Useful utility when users want to compose many transformation.
        - `angles_to_projective_transforms`: Convert angles to transformation matrix.
- tfa/image/translate_ops/translations_to_projective_transforms
    - **Expose**
        - `translations_to_projective_transforms`: Convert translations to transformation matrix.
- tfa/image/utils
    - **Drop**: internal utility
- tfa/losses/metric_learning
    - **Drop**? People can use different distance through `distance` argument for triplet loss.
- tfa/metrics/hamming/hamming_loss_fn
    - **Expose**? I think it's better to name this `hamming_loss` instead of `hamming_loss_fn` for consistency.
- tfa/metrics/utils/sample_weight_shape_match
    - **Drop**: internal utility
- tfa/optimizers/weight_decay_optimizers/DecoupledWeightDecayExtension
    - **Expose**: this is the wrapper of weight decay optimizer.
- tfa/options/warn_fallback
    - **Drop**: internal use.
- tfa/register
    - **Drop**: `tfa.register_all()` has already been exposed.
- tfa/seq2seq/beam_search_decoder
    - **Drop**: main functions/classes have already been exposed to `tfa.seq2seq`.
- tfa/seq2seq/sampler
    - **Drop**: main functions/classes have already been exposed to `tfa.seq2seq`.
- tfa/text/crf:
    - **Expose**? maybe export `tfa.text.crf.CrfDecodeForwardRnnCell` as `tfa.text.CrfDecodeForwardRnnCell`?
WindQAQ: Hmm, team tagging is not so useful. @seanpmorgan @bhack Can you review the comments above when time allows, I would like to move this forward. Thank you!
bhack: @WindQAQ It is the risk of https://en.m.wikipedia.org/wiki/Bystander_effect :smile:

Can you quickly explain what is this?
WindQAQ: > @WindQAQ It is the risk of https://en.m.wikipedia.org/wiki/Bystander_effect 😄

Haha, it's not a good effect actually 😆

> Can you quickly explain what is this?

Okay, so the problem is orignated in https://github.com/tensorflow/addons/issues/2132. 

In short, decide the public API in https://github.com/tensorflow/addons/pull/2162#issuecomment-696876373

**TL;DR**

Changes:

- Fixes #2132. The exposed namespace aligns to the one of core TF: `tf.types.experiemental.*` vs `tfa.types.*`.
- Use `explicit_package_contents_filter` to filter out meaningless endpoints. This also leads to the fact that only modules/functions/classes imported in `__init__.py` will be generated publicly. https://github.com/tensorflow/addons/pull/2162#issuecomment-696876373 lists all modules that I thought to be dropped/exposed after this PR. I will expose them according to the review in this PR as well.
- Pin to newest `tensorflow_docs` version.
bhack: I think Is ok. I don't know much about CRF and seq2seq.
seanpmorgan: > @tensorflow/sig-addons-maintainers I have listed all modules to be discussed down below. Any comment on this? Thank you!
> 
> * tfa/image/color_ops
>   
>   * **Drop**: `equalize` and `sharpness` already exported.
> * image/filters/FloatTensorLike
>   
>   * **Drop**: this is just the alias of `tfa.utils.types.*`
> * tfa/image/transform_ops
>   
>   * **Expose**
>     
>     * `compose_transforms`: Useful utility when users want to compose many transformation.
>     * `angles_to_projective_transforms`: Convert angles to transformation matrix.
> * tfa/image/translate_ops/translations_to_projective_transforms
>   
>   * **Expose**
>     
>     * `translations_to_projective_transforms`: Convert translations to transformation matrix.
> * tfa/image/utils
>   
>   * **Drop**: internal utility
> * tfa/losses/metric_learning
>   
>   * **Drop**? People can use different distance through `distance` argument for triplet loss.
> * tfa/metrics/hamming/hamming_loss_fn
>   
>   * **Expose**? I think it's better to name this `hamming_loss` instead of `hamming_loss_fn` for consistency.
> * tfa/metrics/utils/sample_weight_shape_match
>   
>   * **Drop**: internal utility
> * tfa/optimizers/weight_decay_optimizers/DecoupledWeightDecayExtension
>   
>   * **Expose**: this is the wrapper of weight decay optimizer.
> * tfa/options/warn_fallback
>   
>   * **Drop**: internal use.
> * tfa/register
>   
>   * **Drop**: `tfa.register_all()` has already been exposed.
> * tfa/seq2seq/beam_search_decoder
>   
>   * **Drop**: main functions/classes have already been exposed to `tfa.seq2seq`.
> * tfa/seq2seq/sampler
>   
>   * **Drop**: main functions/classes have already been exposed to `tfa.seq2seq`.
> * tfa/text/crf:
>   
>   * **Expose**? maybe export `tfa.text.crf.CrfDecodeForwardRnnCell` as `tfa.text.CrfDecodeForwardRnnCell`?

Agree with all of the above.... only thing I would change would be to expose all public functions in text.crf as I'm not sure of the user adoption and understanding there. One other thought:

* Expose hamming_loss_fn and we can update the name in another PR (likely need to alias the old name though for backwards compat)

Sorry for the delay forgot to get around to this.

seanpmorgan: @WindQAQ just checking do you think you'll be able to clean up these or would you like someone to pick it up from here? No major rush but was thinking we should get this fixed up prior to Addons pinned to TF2.4
WindQAQ: @seanpmorgan Sure, I'll do it in the weekend. Sorry that I'm busy with my course recently 😭 
WindQAQ: Hi @seanpmorgan @bhack, I've updated the PR. The updated API is here [tfa_api.zip](https://github.com/tensorflow/addons/files/5580808/tfa_api.zip).

Let me now if there is any change needed!


Patch:
diff --git a/tensorflow_addons/__init__.py b/tensorflow_addons/__init__.py
index 2afa5c1..b08f82b 100644
--- a/tensorflow_addons/__init__.py
+++ b/tensorflow_addons/__init__.py
@@ -29,5 +29,6 @@ from tensorflow_addons import rnn
 from tensorflow_addons import seq2seq
 from tensorflow_addons import text
 from tensorflow_addons.register import register_all
+from tensorflow_addons.utils import types
 
 from tensorflow_addons.version import __version__
diff --git a/tensorflow_addons/image/__init__.py b/tensorflow_addons/image/__init__.py
index 5351a26..5781f85 100644
--- a/tensorflow_addons/image/__init__.py
+++ b/tensorflow_addons/image/__init__.py
@@ -34,6 +34,9 @@ from tensorflow_addons.image.transform_ops import rotate
 from tensorflow_addons.image.transform_ops import shear_x
 from tensorflow_addons.image.transform_ops import shear_y
 from tensorflow_addons.image.sparse_image_warp import sparse_image_warp
+from tensorflow_addons.image.transform_ops import compose_transforms
+from tensorflow_addons.image.transform_ops import angles_to_projective_transforms
 from tensorflow_addons.image.transform_ops import transform
 from tensorflow_addons.image.translate_ops import translate
 from tensorflow_addons.image.translate_ops import translate_xy
+from tensorflow_addons.image.translate_ops import translations_to_projective_transforms
diff --git a/tensorflow_addons/metrics/__init__.py b/tensorflow_addons/metrics/__init__.py
index 16df926..8087413 100755
--- a/tensorflow_addons/metrics/__init__.py
+++ b/tensorflow_addons/metrics/__init__.py
@@ -16,7 +16,11 @@
 
 from tensorflow_addons.metrics.cohens_kappa import CohenKappa
 from tensorflow_addons.metrics.f_scores import F1Score, FBetaScore
-from tensorflow_addons.metrics.hamming import HammingLoss, hamming_distance
+from tensorflow_addons.metrics.hamming import (
+    HammingLoss,
+    hamming_distance,
+    hamming_loss_fn,
+)
 from tensorflow_addons.metrics.utils import MeanMetricWrapper
 from tensorflow_addons.metrics.matthews_correlation_coefficient import (
     MatthewsCorrelationCoefficient,
diff --git a/tensorflow_addons/metrics/matthews_correlation_coefficient.py b/tensorflow_addons/metrics/matthews_correlation_coefficient.py
index 0d5a4b9..6600e29 100644
--- a/tensorflow_addons/metrics/matthews_correlation_coefficient.py
+++ b/tensorflow_addons/metrics/matthews_correlation_coefficient.py
@@ -14,6 +14,8 @@
 # ==============================================================================
 """Matthews Correlation Coefficient Implementation."""
 
+import numpy as np
+
 import tensorflow as tf
 from tensorflow.keras import backend as K
 
@@ -148,5 +150,5 @@ class MatthewsCorrelationCoefficient(tf.keras.metrics.Metric):
 
     def reset_states(self):
         """Resets all of the metric state variables."""
-        reset_value = tf.zeros(self.num_classes, dtype=self.dtype)
+        reset_value = np.zeros(self.num_classes, dtype=self.dtype)
         K.batch_set_value([(v, reset_value) for v in self.variables])
diff --git a/tensorflow_addons/optimizers/__init__.py b/tensorflow_addons/optimizers/__init__.py
index a565533..5d03b6f 100644
--- a/tensorflow_addons/optimizers/__init__.py
+++ b/tensorflow_addons/optimizers/__init__.py
@@ -42,4 +42,7 @@ from tensorflow_addons.optimizers.weight_decay_optimizers import SGDW
 from tensorflow_addons.optimizers.weight_decay_optimizers import (
     extend_with_decoupled_weight_decay,
 )
+from tensorflow_addons.optimizers.weight_decay_optimizers import (
+    DecoupledWeightDecayExtension,
+)
 from tensorflow_addons.optimizers.yogi import Yogi
diff --git a/tensorflow_addons/optimizers/average_wrapper.py b/tensorflow_addons/optimizers/average_wrapper.py
index d7c46d2..50b6152 100644
--- a/tensorflow_addons/optimizers/average_wrapper.py
+++ b/tensorflow_addons/optimizers/average_wrapper.py
@@ -18,19 +18,13 @@ import abc
 import tensorflow as tf
 from tensorflow_addons.utils import types
 
-import warnings
 from typeguard import typechecked
-from typing import Optional
 
 
 class AveragedOptimizerWrapper(tf.keras.optimizers.Optimizer, metaclass=abc.ABCMeta):
     @typechecked
     def __init__(
-        self,
-        optimizer: types.Optimizer,
-        sequential_update: Optional[bool] = None,
-        name: str = "AverageOptimizer",
-        **kwargs
+        self, optimizer: types.Optimizer, name: str = "AverageOptimizer", **kwargs
     ):
         super().__init__(name, **kwargs)
 
@@ -42,19 +36,9 @@ class AveragedOptimizerWrapper(tf.keras.optimizers.Optimizer, metaclass=abc.ABCM
                 "optimizer is not an object of tf.keras.optimizers.Optimizer"
             )
 
-        if not isinstance(sequential_update, bool):
-            raise TypeError("sequential_update must be of bool type")
-
         self._optimizer = optimizer
         self._track_trackable(self._optimizer, "awg_optimizer")
 
-        if sequential_update is not None:
-            warnings.warn(
-                "The parameter `sequential_update` is redundant due to AutoGraph. "
-                "This behavior is deprecated and in Addons 0.12, this will raise an error. ",
-                DeprecationWarning,
-            )
-
     def _create_slots(self, var_list):
         self._optimizer._create_slots(var_list=var_list)
         for var in var_list:
diff --git a/tensorflow_addons/optimizers/moving_average.py b/tensorflow_addons/optimizers/moving_average.py
index a46e358..8000e3e 100644
--- a/tensorflow_addons/optimizers/moving_average.py
+++ b/tensorflow_addons/optimizers/moving_average.py
@@ -45,7 +45,6 @@ class MovingAverage(AveragedOptimizerWrapper):
     def __init__(
         self,
         optimizer: types.Optimizer,
-        sequential_update: bool = True,
         average_decay: types.FloatTensorLike = 0.99,
         num_updates: Union[None, int, tf.Variable] = None,
         start_step: int = 0,
@@ -58,10 +57,6 @@ class MovingAverage(AveragedOptimizerWrapper):
         Args:
             optimizer: str or `tf.keras.optimizers.Optimizer` that will be
                 used to compute and apply gradients.
-            sequential_update: Bool. If False, will compute the moving average
-                at the same time as the model is updated, potentially doing
-                benign data races. If True, will update the moving average
-                after gradient updates.
             average_decay: float. Decay to use to maintain the moving averages
                 of trained variables.
             num_updates: Optional count of the number of updates applied to
@@ -79,7 +74,7 @@ class MovingAverage(AveragedOptimizerWrapper):
                 decay of learning rate. `lr` is included for backward
                 compatibility, recommended to use `learning_rate` instead.
         """
-        super().__init__(optimizer, sequential_update, name, **kwargs)
+        super().__init__(optimizer, name, **kwargs)
         self._num_updates = num_updates
         if self._num_updates is not None:
             if isinstance(self._num_updates, tf.Variable):
diff --git a/tensorflow_addons/optimizers/stochastic_weight_averaging.py b/tensorflow_addons/optimizers/stochastic_weight_averaging.py
index adb3df5..cd393f2 100644
--- a/tensorflow_addons/optimizers/stochastic_weight_averaging.py
+++ b/tensorflow_addons/optimizers/stochastic_weight_averaging.py
@@ -78,7 +78,6 @@ class SWA(AveragedOptimizerWrapper):
         start_averaging: int = 0,
         average_period: int = 10,
         name: str = "SWA",
-        sequential_update: bool = True,
         **kwargs
     ):
         r"""Wrap optimizer with the Stochastic Weight Averaging mechanism.
@@ -96,10 +95,6 @@ class SWA(AveragedOptimizerWrapper):
                 needs to be >= 1.
             name: Optional name for the operations created when applying
                 gradients. Defaults to 'SWA'.
-            sequential_update: Bool. If False, will compute the moving average
-                at the same time as the model is updated, potentially doing
-                benign data races. If True, will update the moving average
-                after gradient updates.
             **kwargs: keyword arguments. Allowed to be {`clipnorm`,
                 `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients by
                 norm; `clipvalue` is clip gradients by value, `decay` is
@@ -107,7 +102,7 @@ class SWA(AveragedOptimizerWrapper):
                 decay of learning rate. `lr` is included for backward
                 compatibility, recommended to use `learning_rate` instead.
         """
-        super().__init__(optimizer, sequential_update, name, **kwargs)
+        super().__init__(optimizer, name, **kwargs)
 
         if average_period < 1:
             raise ValueError("average_period must be >= 1")
diff --git a/tensorflow_addons/text/__init__.py b/tensorflow_addons/text/__init__.py
index dfb6379..7f3a890 100644
--- a/tensorflow_addons/text/__init__.py
+++ b/tensorflow_addons/text/__init__.py
@@ -15,6 +15,8 @@
 """Additional text-processing ops."""
 
 # Conditional Random Field
+from tensorflow_addons.text import crf
+from tensorflow_addons.text.crf import CrfDecodeForwardRnnCell
 from tensorflow_addons.text.crf import crf_binary_score
 from tensorflow_addons.text.crf import crf_constrained_decode
 from tensorflow_addons.text.crf import crf_decode
diff --git a/tools/docs/build_docs.py b/tools/docs/build_docs.py
index 4e55502..5442a42 100644
--- a/tools/docs/build_docs.py
+++ b/tools/docs/build_docs.py
@@ -74,8 +74,11 @@ def main(argv):
         py_modules=[(PROJECT_SHORT_NAME, tfa)],
         code_url_prefix=code_url_prefix,
         private_map={"tfa": ["__version__", "utils", "version"]},
-        # This callback usually cleans up a lot of aliases caused by internal imports.
-        callbacks=[public_api.local_definitions_filter],
+        # These callbacks usually clean up a lot of aliases caused by internal imports.
+        callbacks=[
+            public_api.local_definitions_filter,
+            public_api.explicit_package_contents_filter,
+        ],
         search_hints=FLAGS.search_hints,
         site_path=FLAGS.site_path,
     )
diff --git a/tools/install_deps/doc_requirements.txt b/tools/install_deps/doc_requirements.txt
index 0d16271..7a7a9ed 100644
--- a/tools/install_deps/doc_requirements.txt
+++ b/tools/install_deps/doc_requirements.txt
@@ -1,2 +1,2 @@
-git+https://github.com/tensorflow/docs@8192df1
+git+https://github.com/tensorflow/docs
 pyyaml
\ No newline at end of file


Test patch:
diff --git a/tensorflow_addons/callbacks/tests/avg_model_checkpoint_test.py b/tensorflow_addons/callbacks/tests/avg_model_checkpoint_test.py
index 5b6c397..4dbcf2c 100644
--- a/tensorflow_addons/callbacks/tests/avg_model_checkpoint_test.py
+++ b/tensorflow_addons/callbacks/tests/avg_model_checkpoint_test.py
@@ -16,9 +16,7 @@ EPOCHS = 5
 def get_data_and_model(optimizer="moving_avg"):
     x = tf.random.normal([TRAIN_SAMPLES, INPUT_DIM])
     y = tf.random.normal([TRAIN_SAMPLES, NUM_CLASSES])
-    moving_avg = MovingAverage(
-        tf.keras.optimizers.SGD(lr=2.0), sequential_update=True, average_decay=0.5
-    )
+    moving_avg = MovingAverage(tf.keras.optimizers.SGD(lr=2.0), average_decay=0.5)
     if optimizer == "moving_avg":
         optimizer = moving_avg
     inputs = keras.layers.Input(INPUT_DIM)
diff --git a/tensorflow_addons/metrics/tests/matthews_correlation_coefficient_test.py b/tensorflow_addons/metrics/tests/matthews_correlation_coefficient_test.py
index a41f2a9..4f570ad 100644
--- a/tensorflow_addons/metrics/tests/matthews_correlation_coefficient_test.py
+++ b/tensorflow_addons/metrics/tests/matthews_correlation_coefficient_test.py
@@ -14,7 +14,6 @@
 # ==============================================================================
 """Matthews Correlation Coefficient Test."""
 
-
 import tensorflow as tf
 
 import numpy as np
@@ -78,3 +77,18 @@ def test_keras_model():
     labels = np.random.random((10, 1))
     labels = np.where(labels > 0.5, 1.0, 0.0)
     model.fit(data, labels, epochs=1, batch_size=32, verbose=0)
+
+
+def test_reset_states_graph():
+    gt_label = tf.constant([[1.0], [1.0], [1.0], [0.0]], dtype=tf.float32)
+    preds = tf.constant([[1.0], [0.0], [1.0], [1.0]], dtype=tf.float32)
+    mcc = MatthewsCorrelationCoefficient(1)
+    mcc.update_state(gt_label, preds)
+
+    @tf.function
+    def reset_states():
+        mcc.reset_states()
+
+    reset_states()
+    # Check results
+    check_results(mcc, [0])