diff --git a/tensorflow_ranking/g3doc/api_docs/python/index.md b/tensorflow_ranking/g3doc/api_docs/python/index.md
deleted file mode 100644
index 4a4c760..0000000
--- a/tensorflow_ranking/g3doc/api_docs/python/index.md
+++ /dev/null
@@ -1,117 +0,0 @@
-# All symbols in TensorFlow Ranking
-
-<!-- Insert buttons and diff -->
-
-## Primary symbols
-
-*   <a href="./tfr.md"><code>tfr</code></a>
-*   <a href="./tfr/data.md"><code>tfr.data</code></a>
-*   <a href="./tfr/data/build_ranking_dataset.md"><code>tfr.data.build_ranking_dataset</code></a>
-*   <a href="./tfr/data/build_ranking_dataset_with_parsing_fn.md"><code>tfr.data.build_ranking_dataset_with_parsing_fn</code></a>
-*   <a href="./tfr/data/build_ranking_serving_input_receiver_fn.md"><code>tfr.data.build_ranking_serving_input_receiver_fn</code></a>
-*   <a href="./tfr/data/build_ranking_serving_input_receiver_fn_with_parsing_fn.md"><code>tfr.data.build_ranking_serving_input_receiver_fn_with_parsing_fn</code></a>
-*   <a href="./tfr/data/build_sequence_example_serving_input_receiver_fn.md"><code>tfr.data.build_sequence_example_serving_input_receiver_fn</code></a>
-*   <a href="./tfr/data/make_parsing_fn.md"><code>tfr.data.make_parsing_fn</code></a>
-*   <a href="./tfr/data/parse_from_example_in_example.md"><code>tfr.data.parse_from_example_in_example</code></a>
-*   <a href="./tfr/data/parse_from_example_list.md"><code>tfr.data.parse_from_example_list</code></a>
-*   <a href="./tfr/data/parse_from_sequence_example.md"><code>tfr.data.parse_from_sequence_example</code></a>
-*   <a href="./tfr/data/read_batched_sequence_example_dataset.md"><code>tfr.data.read_batched_sequence_example_dataset</code></a>
-*   <a href="./tfr/estimator.md"><code>tfr.estimator</code></a>
-*   <a href="./tfr/estimator/EstimatorBuilder.md"><code>tfr.estimator.EstimatorBuilder</code></a>
-*   <a href="./tfr/estimator/GAMEstimatorBuilder.md"><code>tfr.estimator.GAMEstimatorBuilder</code></a>
-*   <a href="./tfr/estimator/make_dnn_ranking_estimator.md"><code>tfr.estimator.make_dnn_ranking_estimator</code></a>
-*   <a href="./tfr/estimator/make_gam_ranking_estimator.md"><code>tfr.estimator.make_gam_ranking_estimator</code></a>
-*   <a href="./tfr/ext.md"><code>tfr.ext</code></a>
-*   <a href="./tfr/ext/pipeline.md"><code>tfr.ext.pipeline</code></a>
-*   <a href="./tfr/ext/pipeline/RankingPipeline.md"><code>tfr.ext.pipeline.RankingPipeline</code></a>
-*   <a href="./tfr/ext/tfrbert.md"><code>tfr.ext.tfrbert</code></a>
-*   <a href="./tfr/ext/tfrbert/TFRBertRankingNetwork.md"><code>tfr.ext.tfrbert.TFRBertRankingNetwork</code></a>
-*   <a href="./tfr/ext/tfrbert/TFRBertUtil.md"><code>tfr.ext.tfrbert.TFRBertUtil</code></a>
-*   <a href="./tfr/feature.md"><code>tfr.feature</code></a>
-*   <a href="./tfr/feature/encode_features.md"><code>tfr.feature.encode_features</code></a>
-*   <a href="./tfr/feature/encode_listwise_features.md"><code>tfr.feature.encode_listwise_features</code></a>
-*   <a href="./tfr/feature/encode_pointwise_features.md"><code>tfr.feature.encode_pointwise_features</code></a>
-*   <a href="./tfr/feature/make_identity_transform_fn.md"><code>tfr.feature.make_identity_transform_fn</code></a>
-*   <a href="./tfr/head.md"><code>tfr.head</code></a>
-*   <a href="./tfr/head/create_multi_ranking_head.md"><code>tfr.head.create_multi_ranking_head</code></a>
-*   <a href="./tfr/head/create_ranking_head.md"><code>tfr.head.create_ranking_head</code></a>
-*   <a href="./tfr/keras.md"><code>tfr.keras</code></a>
-*   <a href="./tfr/keras/canned.md"><code>tfr.keras.canned</code></a>
-*   <a href="./tfr/keras/canned/DNNRankingNetwork.md"><code>tfr.keras.canned.DNNRankingNetwork</code></a>
-*   <a href="./tfr/keras/canned/GAMRankingNetwork.md"><code>tfr.keras.canned.GAMRankingNetwork</code></a>
-*   <a href="./tfr/keras/canned/dnn.md"><code>tfr.keras.canned.dnn</code></a>
-*   <a href="./tfr/keras/canned/DNNRankingNetwork.md"><code>tfr.keras.canned.dnn.DNNRankingNetwork</code></a>
-*   <a href="./tfr/keras/canned/gam.md"><code>tfr.keras.canned.gam</code></a>
-*   <a href="./tfr/keras/canned/GAMRankingNetwork.md"><code>tfr.keras.canned.gam.GAMRankingNetwork</code></a>
-*   <a href="./tfr/keras/estimator.md"><code>tfr.keras.estimator</code></a>
-*   <a href="./tfr/keras/estimator/model_to_estimator.md"><code>tfr.keras.estimator.model_to_estimator</code></a>
-*   <a href="./tfr/keras/feature.md"><code>tfr.keras.feature</code></a>
-*   <a href="./tfr/keras/feature/EncodeListwiseFeatures.md"><code>tfr.keras.feature.EncodeListwiseFeatures</code></a>
-*   <a href="./tfr/keras/feature/GenerateMask.md"><code>tfr.keras.feature.GenerateMask</code></a>
-*   <a href="./tfr/keras/feature/create_keras_inputs.md"><code>tfr.keras.feature.create_keras_inputs</code></a>
-*   <a href="./tfr/keras/feature/deserialize_feature_columns.md"><code>tfr.keras.feature.deserialize_feature_columns</code></a>
-*   <a href="./tfr/keras/feature/serialize_feature_columns.md"><code>tfr.keras.feature.serialize_feature_columns</code></a>
-*   <a href="./tfr/keras/losses.md"><code>tfr.keras.losses</code></a>
-*   <a href="./tfr/keras/losses/ApproxMRRLoss.md"><code>tfr.keras.losses.ApproxMRRLoss</code></a>
-*   <a href="./tfr/keras/losses/ApproxNDCGLoss.md"><code>tfr.keras.losses.ApproxNDCGLoss</code></a>
-*   <a href="./tfr/keras/losses/GumbelApproxNDCGLoss.md"><code>tfr.keras.losses.GumbelApproxNDCGLoss</code></a>
-*   <a href="./tfr/keras/losses/ListMLELoss.md"><code>tfr.keras.losses.ListMLELoss</code></a>
-*   <a href="./tfr/keras/losses/MeanSquaredLoss.md"><code>tfr.keras.losses.MeanSquaredLoss</code></a>
-*   <a href="./tfr/keras/losses/PairwiseHingeLoss.md"><code>tfr.keras.losses.PairwiseHingeLoss</code></a>
-*   <a href="./tfr/keras/losses/PairwiseLogisticLoss.md"><code>tfr.keras.losses.PairwiseLogisticLoss</code></a>
-*   <a href="./tfr/keras/losses/PairwiseSoftZeroOneLoss.md"><code>tfr.keras.losses.PairwiseSoftZeroOneLoss</code></a>
-*   <a href="./tfr/keras/losses/RankingLossKey.md"><code>tfr.keras.losses.RankingLossKey</code></a>
-*   <a href="./tfr/keras/losses/SigmoidCrossEntropyLoss.md"><code>tfr.keras.losses.SigmoidCrossEntropyLoss</code></a>
-*   <a href="./tfr/keras/losses/SoftmaxLoss.md"><code>tfr.keras.losses.SoftmaxLoss</code></a>
-*   <a href="./tfr/keras/losses/UniqueSoftmaxLoss.md"><code>tfr.keras.losses.UniqueSoftmaxLoss</code></a>
-*   <a href="./tfr/keras/losses/get.md"><code>tfr.keras.losses.get</code></a>
-*   <a href="./tfr/keras/metrics.md"><code>tfr.keras.metrics</code></a>
-*   <a href="./tfr/keras/metrics/ARPMetric.md"><code>tfr.keras.metrics.ARPMetric</code></a>
-*   <a href="./tfr/keras/metrics/AlphaDCGMetric.md"><code>tfr.keras.metrics.AlphaDCGMetric</code></a>
-*   <a href="./tfr/keras/metrics/DCGMetric.md"><code>tfr.keras.metrics.DCGMetric</code></a>
-*   <a href="./tfr/keras/metrics/MRRMetric.md"><code>tfr.keras.metrics.MRRMetric</code></a>
-*   <a href="./tfr/keras/metrics/MeanAveragePrecisionMetric.md"><code>tfr.keras.metrics.MeanAveragePrecisionMetric</code></a>
-*   <a href="./tfr/keras/metrics/NDCGMetric.md"><code>tfr.keras.metrics.NDCGMetric</code></a>
-*   <a href="./tfr/keras/metrics/OPAMetric.md"><code>tfr.keras.metrics.OPAMetric</code></a>
-*   <a href="./tfr/keras/metrics/PrecisionIAMetric.md"><code>tfr.keras.metrics.PrecisionIAMetric</code></a>
-*   <a href="./tfr/keras/metrics/PrecisionMetric.md"><code>tfr.keras.metrics.PrecisionMetric</code></a>
-*   <a href="./tfr/keras/metrics/default_keras_metrics.md"><code>tfr.keras.metrics.default_keras_metrics</code></a>
-*   <a href="./tfr/keras/model.md"><code>tfr.keras.model</code></a>
-*   <a href="./tfr/keras/model/create_keras_model.md"><code>tfr.keras.model.create_keras_model</code></a>
-*   <a href="./tfr/keras/network.md"><code>tfr.keras.network</code></a>
-*   <a href="./tfr/keras/network/RankingNetwork.md"><code>tfr.keras.network.RankingNetwork</code></a>
-*   <a href="./tfr/keras/network/UnivariateRankingNetwork.md"><code>tfr.keras.network.UnivariateRankingNetwork</code></a>
-*   <a href="./tfr/keras/network/listwise_scoring.md"><code>tfr.keras.network.listwise_scoring</code></a>
-*   <a href="./tfr/losses.md"><code>tfr.losses</code></a>
-*   <a href="./tfr/losses/RankingLossKey.md"><code>tfr.losses.RankingLossKey</code></a>
-*   <a href="./tfr/losses/create_ndcg_lambda_weight.md"><code>tfr.losses.create_ndcg_lambda_weight</code></a>
-*   <a href="./tfr/losses/create_p_list_mle_lambda_weight.md"><code>tfr.losses.create_p_list_mle_lambda_weight</code></a>
-*   <a href="./tfr/losses/create_reciprocal_rank_lambda_weight.md"><code>tfr.losses.create_reciprocal_rank_lambda_weight</code></a>
-*   <a href="./tfr/losses/make_loss_fn.md"><code>tfr.losses.make_loss_fn</code></a>
-*   <a href="./tfr/losses/make_loss_metric_fn.md"><code>tfr.losses.make_loss_metric_fn</code></a>
-*   <a href="./tfr/metrics.md"><code>tfr.metrics</code></a>
-*   <a href="./tfr/metrics/RankingMetricKey.md"><code>tfr.metrics.RankingMetricKey</code></a>
-*   <a href="./tfr/metrics/alpha_discounted_cumulative_gain.md"><code>tfr.metrics.alpha_discounted_cumulative_gain</code></a>
-*   <a href="./tfr/metrics/average_relevance_position.md"><code>tfr.metrics.average_relevance_position</code></a>
-*   <a href="./tfr/metrics/compute_mean.md"><code>tfr.metrics.compute_mean</code></a>
-*   <a href="./tfr/metrics/discounted_cumulative_gain.md"><code>tfr.metrics.discounted_cumulative_gain</code></a>
-*   <a href="./tfr/metrics/eval_metric.md"><code>tfr.metrics.eval_metric</code></a>
-*   <a href="./tfr/metrics/make_ranking_metric_fn.md"><code>tfr.metrics.make_ranking_metric_fn</code></a>
-*   <a href="./tfr/metrics/mean_average_precision.md"><code>tfr.metrics.mean_average_precision</code></a>
-*   <a href="./tfr/metrics/mean_reciprocal_rank.md"><code>tfr.metrics.mean_reciprocal_rank</code></a>
-*   <a href="./tfr/metrics/normalized_discounted_cumulative_gain.md"><code>tfr.metrics.normalized_discounted_cumulative_gain</code></a>
-*   <a href="./tfr/metrics/ordered_pair_accuracy.md"><code>tfr.metrics.ordered_pair_accuracy</code></a>
-*   <a href="./tfr/metrics/precision.md"><code>tfr.metrics.precision</code></a>
-*   <a href="./tfr/metrics/precision_ia.md"><code>tfr.metrics.precision_ia</code></a>
-*   <a href="./tfr/model.md"><code>tfr.model</code></a>
-*   <a href="./tfr/model/make_groupwise_ranking_fn.md"><code>tfr.model.make_groupwise_ranking_fn</code></a>
-*   <a href="./tfr/utils.md"><code>tfr.utils</code></a>
-*   <a href="./tfr/utils/de_noise.md"><code>tfr.utils.de_noise</code></a>
-*   <a href="./tfr/utils/is_label_valid.md"><code>tfr.utils.is_label_valid</code></a>
-*   <a href="./tfr/utils/organize_valid_indices.md"><code>tfr.utils.organize_valid_indices</code></a>
-*   <a href="./tfr/utils/padded_nd_indices.md"><code>tfr.utils.padded_nd_indices</code></a>
-*   <a href="./tfr/utils/reshape_first_ndims.md"><code>tfr.utils.reshape_first_ndims</code></a>
-*   <a href="./tfr/utils/reshape_to_2d.md"><code>tfr.utils.reshape_to_2d</code></a>
-*   <a href="./tfr/utils/shuffle_valid_indices.md"><code>tfr.utils.shuffle_valid_indices</code></a>
-*   <a href="./tfr/utils/sort_by_scores.md"><code>tfr.utils.sort_by_scores</code></a>
-*   <a href="./tfr/utils/sorted_ranks.md"><code>tfr.utils.sorted_ranks</code></a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr.md b/tensorflow_ranking/g3doc/api_docs/python/tfr.md
index c5932bd..445ad34 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr.md
@@ -1,3 +1,6 @@
+description: TensorFlow Ranking, the project to build ranking models on
+TensorFlow.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr" />
 <meta itemprop="path" content="Stable" />
@@ -8,8 +11,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/__init__.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
@@ -44,6 +46,18 @@ library.
 
 [`utils`](./tfr/utils.md) module: Utility functions for ranking library.
 
-## Other Members
+<!-- Tabular view -->
 
-*   `__version__ = '0.3.2.dev'` <a id="__version__"></a>
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Other Members</h2></th></tr>
+
+<tr>
+<td>
+__version__<a id="__version__"></a>
+</td>
+<td>
+`'0.3.3.dev'`
+</td>
+</tr>
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/_api_cache.json b/tensorflow_ranking/g3doc/api_docs/python/tfr/_api_cache.json
index 0171cb6..e4830de 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/_api_cache.json
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/_api_cache.json
@@ -1,8 +1,8 @@
 {
   "duplicate_of": {
-    "tfr.data.absolute_import": "tfr.keras.losses.absolute_import",
-    "tfr.data.division": "tfr.keras.losses.division",
-    "tfr.data.print_function": "tfr.keras.losses.print_function",
+    "tfr.data.absolute_import": "tfr.keras.metrics.absolute_import",
+    "tfr.data.division": "tfr.keras.metrics.division",
+    "tfr.data.print_function": "tfr.keras.metrics.print_function",
     "tfr.estimator.EstimatorBuilder.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.estimator.EstimatorBuilder.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
     "tfr.estimator.EstimatorBuilder.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
@@ -19,9 +19,9 @@
     "tfr.estimator.GAMEstimatorBuilder.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
     "tfr.estimator.GAMEstimatorBuilder.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
     "tfr.estimator.GAMEstimatorBuilder.make_estimator": "tfr.estimator.EstimatorBuilder.make_estimator",
-    "tfr.estimator.absolute_import": "tfr.keras.losses.absolute_import",
-    "tfr.estimator.division": "tfr.keras.losses.division",
-    "tfr.estimator.print_function": "tfr.keras.losses.print_function",
+    "tfr.estimator.absolute_import": "tfr.keras.metrics.absolute_import",
+    "tfr.estimator.division": "tfr.keras.metrics.division",
+    "tfr.estimator.print_function": "tfr.keras.metrics.print_function",
     "tfr.ext.pipeline.RankingPipeline.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.ext.pipeline.RankingPipeline.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
     "tfr.ext.pipeline.RankingPipeline.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
@@ -29,9 +29,9 @@
     "tfr.ext.pipeline.RankingPipeline.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
     "tfr.ext.pipeline.RankingPipeline.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
     "tfr.ext.pipeline.RankingPipeline.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
-    "tfr.ext.pipeline.absolute_import": "tfr.keras.losses.absolute_import",
-    "tfr.ext.pipeline.division": "tfr.keras.losses.division",
-    "tfr.ext.pipeline.print_function": "tfr.keras.losses.print_function",
+    "tfr.ext.pipeline.absolute_import": "tfr.keras.metrics.absolute_import",
+    "tfr.ext.pipeline.division": "tfr.keras.metrics.division",
+    "tfr.ext.pipeline.print_function": "tfr.keras.metrics.print_function",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.__call__": "tfr.keras.canned.DNNRankingNetwork.__call__",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
@@ -46,6 +46,7 @@
     "tfr.ext.tfrbert.TFRBertRankingNetwork.add_weight": "tfr.keras.canned.DNNRankingNetwork.add_weight",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.build": "tfr.keras.canned.DNNRankingNetwork.build",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.call": "tfr.keras.network.RankingNetwork.call",
+    "tfr.ext.tfrbert.TFRBertRankingNetwork.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.compute_logits": "tfr.keras.network.UnivariateRankingNetwork.compute_logits",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
@@ -53,6 +54,7 @@
     "tfr.ext.tfrbert.TFRBertRankingNetwork.context_feature_columns": "tfr.keras.network.RankingNetwork.context_feature_columns",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.dtype": "tfr.keras.canned.DNNRankingNetwork.dtype",
+    "tfr.ext.tfrbert.TFRBertRankingNetwork.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.example_feature_columns": "tfr.keras.network.RankingNetwork.example_feature_columns",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
@@ -70,6 +72,7 @@
     "tfr.ext.tfrbert.TFRBertRankingNetwork.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.transform": "tfr.keras.network.RankingNetwork.transform",
+    "tfr.ext.tfrbert.TFRBertRankingNetwork.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.ext.tfrbert.TFRBertRankingNetwork.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.ext.tfrbert.TFRBertUtil.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.ext.tfrbert.TFRBertUtil.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
@@ -78,15 +81,15 @@
     "tfr.ext.tfrbert.TFRBertUtil.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
     "tfr.ext.tfrbert.TFRBertUtil.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
     "tfr.ext.tfrbert.TFRBertUtil.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
-    "tfr.ext.tfrbert.absolute_import": "tfr.keras.losses.absolute_import",
-    "tfr.ext.tfrbert.division": "tfr.keras.losses.division",
-    "tfr.ext.tfrbert.print_function": "tfr.keras.losses.print_function",
-    "tfr.feature.absolute_import": "tfr.keras.losses.absolute_import",
-    "tfr.feature.division": "tfr.keras.losses.division",
-    "tfr.feature.print_function": "tfr.keras.losses.print_function",
-    "tfr.head.absolute_import": "tfr.keras.losses.absolute_import",
-    "tfr.head.division": "tfr.keras.losses.division",
-    "tfr.head.print_function": "tfr.keras.losses.print_function",
+    "tfr.ext.tfrbert.absolute_import": "tfr.keras.metrics.absolute_import",
+    "tfr.ext.tfrbert.division": "tfr.keras.metrics.division",
+    "tfr.ext.tfrbert.print_function": "tfr.keras.metrics.print_function",
+    "tfr.feature.absolute_import": "tfr.keras.metrics.absolute_import",
+    "tfr.feature.division": "tfr.keras.metrics.division",
+    "tfr.feature.print_function": "tfr.keras.metrics.print_function",
+    "tfr.head.absolute_import": "tfr.keras.metrics.absolute_import",
+    "tfr.head.division": "tfr.keras.metrics.division",
+    "tfr.head.print_function": "tfr.keras.metrics.print_function",
     "tfr.keras.canned.DNNRankingNetwork.call": "tfr.keras.network.RankingNetwork.call",
     "tfr.keras.canned.DNNRankingNetwork.compute_logits": "tfr.keras.network.UnivariateRankingNetwork.compute_logits",
     "tfr.keras.canned.DNNRankingNetwork.context_feature_columns": "tfr.keras.network.RankingNetwork.context_feature_columns",
@@ -106,6 +109,7 @@
     "tfr.keras.canned.GAMRankingNetwork.add_weight": "tfr.keras.canned.DNNRankingNetwork.add_weight",
     "tfr.keras.canned.GAMRankingNetwork.build": "tfr.keras.canned.DNNRankingNetwork.build",
     "tfr.keras.canned.GAMRankingNetwork.call": "tfr.keras.network.RankingNetwork.call",
+    "tfr.keras.canned.GAMRankingNetwork.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.canned.GAMRankingNetwork.compute_logits": "tfr.keras.network.UnivariateRankingNetwork.compute_logits",
     "tfr.keras.canned.GAMRankingNetwork.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.canned.GAMRankingNetwork.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
@@ -113,6 +117,7 @@
     "tfr.keras.canned.GAMRankingNetwork.context_feature_columns": "tfr.keras.network.RankingNetwork.context_feature_columns",
     "tfr.keras.canned.GAMRankingNetwork.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.canned.GAMRankingNetwork.dtype": "tfr.keras.canned.DNNRankingNetwork.dtype",
+    "tfr.keras.canned.GAMRankingNetwork.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.canned.GAMRankingNetwork.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.canned.GAMRankingNetwork.example_feature_columns": "tfr.keras.network.RankingNetwork.example_feature_columns",
     "tfr.keras.canned.GAMRankingNetwork.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
@@ -130,6 +135,7 @@
     "tfr.keras.canned.GAMRankingNetwork.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
     "tfr.keras.canned.GAMRankingNetwork.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
     "tfr.keras.canned.GAMRankingNetwork.transform": "tfr.keras.network.RankingNetwork.transform",
+    "tfr.keras.canned.GAMRankingNetwork.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.canned.GAMRankingNetwork.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.canned.dnn.DNNRankingNetwork": "tfr.keras.canned.DNNRankingNetwork",
     "tfr.keras.canned.dnn.DNNRankingNetwork.__call__": "tfr.keras.canned.DNNRankingNetwork.__call__",
@@ -147,6 +153,7 @@
     "tfr.keras.canned.dnn.DNNRankingNetwork.add_weight": "tfr.keras.canned.DNNRankingNetwork.add_weight",
     "tfr.keras.canned.dnn.DNNRankingNetwork.build": "tfr.keras.canned.DNNRankingNetwork.build",
     "tfr.keras.canned.dnn.DNNRankingNetwork.call": "tfr.keras.network.RankingNetwork.call",
+    "tfr.keras.canned.dnn.DNNRankingNetwork.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.canned.dnn.DNNRankingNetwork.compute_logits": "tfr.keras.network.UnivariateRankingNetwork.compute_logits",
     "tfr.keras.canned.dnn.DNNRankingNetwork.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.canned.dnn.DNNRankingNetwork.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
@@ -154,6 +161,7 @@
     "tfr.keras.canned.dnn.DNNRankingNetwork.context_feature_columns": "tfr.keras.network.RankingNetwork.context_feature_columns",
     "tfr.keras.canned.dnn.DNNRankingNetwork.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.canned.dnn.DNNRankingNetwork.dtype": "tfr.keras.canned.DNNRankingNetwork.dtype",
+    "tfr.keras.canned.dnn.DNNRankingNetwork.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.canned.dnn.DNNRankingNetwork.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.canned.dnn.DNNRankingNetwork.example_feature_columns": "tfr.keras.network.RankingNetwork.example_feature_columns",
     "tfr.keras.canned.dnn.DNNRankingNetwork.get_config": "tfr.keras.canned.DNNRankingNetwork.get_config",
@@ -173,6 +181,7 @@
     "tfr.keras.canned.dnn.DNNRankingNetwork.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
     "tfr.keras.canned.dnn.DNNRankingNetwork.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
     "tfr.keras.canned.dnn.DNNRankingNetwork.transform": "tfr.keras.network.RankingNetwork.transform",
+    "tfr.keras.canned.dnn.DNNRankingNetwork.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.canned.dnn.DNNRankingNetwork.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.canned.gam.GAMRankingNetwork": "tfr.keras.canned.GAMRankingNetwork",
     "tfr.keras.canned.gam.GAMRankingNetwork.__call__": "tfr.keras.canned.DNNRankingNetwork.__call__",
@@ -190,6 +199,7 @@
     "tfr.keras.canned.gam.GAMRankingNetwork.add_weight": "tfr.keras.canned.DNNRankingNetwork.add_weight",
     "tfr.keras.canned.gam.GAMRankingNetwork.build": "tfr.keras.canned.DNNRankingNetwork.build",
     "tfr.keras.canned.gam.GAMRankingNetwork.call": "tfr.keras.network.RankingNetwork.call",
+    "tfr.keras.canned.gam.GAMRankingNetwork.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.canned.gam.GAMRankingNetwork.compute_logits": "tfr.keras.network.UnivariateRankingNetwork.compute_logits",
     "tfr.keras.canned.gam.GAMRankingNetwork.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.canned.gam.GAMRankingNetwork.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
@@ -197,6 +207,7 @@
     "tfr.keras.canned.gam.GAMRankingNetwork.context_feature_columns": "tfr.keras.network.RankingNetwork.context_feature_columns",
     "tfr.keras.canned.gam.GAMRankingNetwork.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.canned.gam.GAMRankingNetwork.dtype": "tfr.keras.canned.DNNRankingNetwork.dtype",
+    "tfr.keras.canned.gam.GAMRankingNetwork.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.canned.gam.GAMRankingNetwork.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.canned.gam.GAMRankingNetwork.example_feature_columns": "tfr.keras.network.RankingNetwork.example_feature_columns",
     "tfr.keras.canned.gam.GAMRankingNetwork.get_config": "tfr.keras.canned.GAMRankingNetwork.get_config",
@@ -216,6 +227,7 @@
     "tfr.keras.canned.gam.GAMRankingNetwork.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
     "tfr.keras.canned.gam.GAMRankingNetwork.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
     "tfr.keras.canned.gam.GAMRankingNetwork.transform": "tfr.keras.network.RankingNetwork.transform",
+    "tfr.keras.canned.gam.GAMRankingNetwork.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.canned.gam.GAMRankingNetwork.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.feature.EncodeListwiseFeatures.__call__": "tfr.keras.canned.DNNRankingNetwork.__call__",
     "tfr.keras.feature.EncodeListwiseFeatures.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
@@ -230,11 +242,13 @@
     "tfr.keras.feature.EncodeListwiseFeatures.add_metric": "tfr.keras.canned.DNNRankingNetwork.add_metric",
     "tfr.keras.feature.EncodeListwiseFeatures.add_weight": "tfr.keras.canned.DNNRankingNetwork.add_weight",
     "tfr.keras.feature.EncodeListwiseFeatures.build": "tfr.keras.canned.DNNRankingNetwork.build",
+    "tfr.keras.feature.EncodeListwiseFeatures.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.feature.EncodeListwiseFeatures.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.feature.EncodeListwiseFeatures.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
     "tfr.keras.feature.EncodeListwiseFeatures.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
     "tfr.keras.feature.EncodeListwiseFeatures.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.feature.EncodeListwiseFeatures.dtype": "tfr.keras.canned.DNNRankingNetwork.dtype",
+    "tfr.keras.feature.EncodeListwiseFeatures.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.feature.EncodeListwiseFeatures.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.feature.EncodeListwiseFeatures.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
     "tfr.keras.feature.EncodeListwiseFeatures.input": "tfr.keras.canned.DNNRankingNetwork.input",
@@ -250,6 +264,7 @@
     "tfr.keras.feature.EncodeListwiseFeatures.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
     "tfr.keras.feature.EncodeListwiseFeatures.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
     "tfr.keras.feature.EncodeListwiseFeatures.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.feature.EncodeListwiseFeatures.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.feature.EncodeListwiseFeatures.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.feature.GenerateMask.__call__": "tfr.keras.canned.DNNRankingNetwork.__call__",
     "tfr.keras.feature.GenerateMask.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
@@ -264,11 +279,13 @@
     "tfr.keras.feature.GenerateMask.add_metric": "tfr.keras.canned.DNNRankingNetwork.add_metric",
     "tfr.keras.feature.GenerateMask.add_weight": "tfr.keras.canned.DNNRankingNetwork.add_weight",
     "tfr.keras.feature.GenerateMask.build": "tfr.keras.canned.DNNRankingNetwork.build",
+    "tfr.keras.feature.GenerateMask.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.feature.GenerateMask.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.feature.GenerateMask.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
     "tfr.keras.feature.GenerateMask.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
     "tfr.keras.feature.GenerateMask.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.feature.GenerateMask.dtype": "tfr.keras.canned.DNNRankingNetwork.dtype",
+    "tfr.keras.feature.GenerateMask.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.feature.GenerateMask.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.feature.GenerateMask.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
     "tfr.keras.feature.GenerateMask.input": "tfr.keras.canned.DNNRankingNetwork.input",
@@ -284,7 +301,156 @@
     "tfr.keras.feature.GenerateMask.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
     "tfr.keras.feature.GenerateMask.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
     "tfr.keras.feature.GenerateMask.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.feature.GenerateMask.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.feature.GenerateMask.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
+    "tfr.keras.layers.ConcatFeatures.__call__": "tfr.keras.canned.DNNRankingNetwork.__call__",
+    "tfr.keras.layers.ConcatFeatures.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
+    "tfr.keras.layers.ConcatFeatures.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
+    "tfr.keras.layers.ConcatFeatures.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
+    "tfr.keras.layers.ConcatFeatures.__le__": "tfr.keras.canned.DNNRankingNetwork.__le__",
+    "tfr.keras.layers.ConcatFeatures.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
+    "tfr.keras.layers.ConcatFeatures.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
+    "tfr.keras.layers.ConcatFeatures.__new__": "tfr.keras.canned.DNNRankingNetwork.__new__",
+    "tfr.keras.layers.ConcatFeatures.activity_regularizer": "tfr.keras.canned.DNNRankingNetwork.activity_regularizer",
+    "tfr.keras.layers.ConcatFeatures.add_loss": "tfr.keras.canned.DNNRankingNetwork.add_loss",
+    "tfr.keras.layers.ConcatFeatures.add_metric": "tfr.keras.canned.DNNRankingNetwork.add_metric",
+    "tfr.keras.layers.ConcatFeatures.add_weight": "tfr.keras.canned.DNNRankingNetwork.add_weight",
+    "tfr.keras.layers.ConcatFeatures.build": "tfr.keras.canned.DNNRankingNetwork.build",
+    "tfr.keras.layers.ConcatFeatures.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
+    "tfr.keras.layers.ConcatFeatures.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
+    "tfr.keras.layers.ConcatFeatures.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
+    "tfr.keras.layers.ConcatFeatures.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
+    "tfr.keras.layers.ConcatFeatures.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
+    "tfr.keras.layers.ConcatFeatures.dtype": "tfr.keras.canned.DNNRankingNetwork.dtype",
+    "tfr.keras.layers.ConcatFeatures.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
+    "tfr.keras.layers.ConcatFeatures.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
+    "tfr.keras.layers.ConcatFeatures.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
+    "tfr.keras.layers.ConcatFeatures.input": "tfr.keras.canned.DNNRankingNetwork.input",
+    "tfr.keras.layers.ConcatFeatures.input_spec": "tfr.keras.canned.DNNRankingNetwork.input_spec",
+    "tfr.keras.layers.ConcatFeatures.losses": "tfr.keras.canned.DNNRankingNetwork.losses",
+    "tfr.keras.layers.ConcatFeatures.metrics": "tfr.keras.canned.DNNRankingNetwork.metrics",
+    "tfr.keras.layers.ConcatFeatures.name": "tfr.keras.canned.DNNRankingNetwork.name",
+    "tfr.keras.layers.ConcatFeatures.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
+    "tfr.keras.layers.ConcatFeatures.non_trainable_weights": "tfr.keras.canned.DNNRankingNetwork.non_trainable_weights",
+    "tfr.keras.layers.ConcatFeatures.output": "tfr.keras.canned.DNNRankingNetwork.output",
+    "tfr.keras.layers.ConcatFeatures.set_weights": "tfr.keras.canned.DNNRankingNetwork.set_weights",
+    "tfr.keras.layers.ConcatFeatures.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
+    "tfr.keras.layers.ConcatFeatures.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
+    "tfr.keras.layers.ConcatFeatures.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
+    "tfr.keras.layers.ConcatFeatures.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.layers.ConcatFeatures.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
+    "tfr.keras.layers.ConcatFeatures.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
+    "tfr.keras.layers.DocumentInteractionAttention.__call__": "tfr.keras.canned.DNNRankingNetwork.__call__",
+    "tfr.keras.layers.DocumentInteractionAttention.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
+    "tfr.keras.layers.DocumentInteractionAttention.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
+    "tfr.keras.layers.DocumentInteractionAttention.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
+    "tfr.keras.layers.DocumentInteractionAttention.__le__": "tfr.keras.canned.DNNRankingNetwork.__le__",
+    "tfr.keras.layers.DocumentInteractionAttention.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
+    "tfr.keras.layers.DocumentInteractionAttention.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
+    "tfr.keras.layers.DocumentInteractionAttention.__new__": "tfr.keras.canned.DNNRankingNetwork.__new__",
+    "tfr.keras.layers.DocumentInteractionAttention.activity_regularizer": "tfr.keras.canned.DNNRankingNetwork.activity_regularizer",
+    "tfr.keras.layers.DocumentInteractionAttention.add_loss": "tfr.keras.canned.DNNRankingNetwork.add_loss",
+    "tfr.keras.layers.DocumentInteractionAttention.add_metric": "tfr.keras.canned.DNNRankingNetwork.add_metric",
+    "tfr.keras.layers.DocumentInteractionAttention.add_weight": "tfr.keras.canned.DNNRankingNetwork.add_weight",
+    "tfr.keras.layers.DocumentInteractionAttention.build": "tfr.keras.canned.DNNRankingNetwork.build",
+    "tfr.keras.layers.DocumentInteractionAttention.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
+    "tfr.keras.layers.DocumentInteractionAttention.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
+    "tfr.keras.layers.DocumentInteractionAttention.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
+    "tfr.keras.layers.DocumentInteractionAttention.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
+    "tfr.keras.layers.DocumentInteractionAttention.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
+    "tfr.keras.layers.DocumentInteractionAttention.dtype": "tfr.keras.canned.DNNRankingNetwork.dtype",
+    "tfr.keras.layers.DocumentInteractionAttention.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
+    "tfr.keras.layers.DocumentInteractionAttention.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
+    "tfr.keras.layers.DocumentInteractionAttention.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
+    "tfr.keras.layers.DocumentInteractionAttention.input": "tfr.keras.canned.DNNRankingNetwork.input",
+    "tfr.keras.layers.DocumentInteractionAttention.input_spec": "tfr.keras.canned.DNNRankingNetwork.input_spec",
+    "tfr.keras.layers.DocumentInteractionAttention.losses": "tfr.keras.canned.DNNRankingNetwork.losses",
+    "tfr.keras.layers.DocumentInteractionAttention.metrics": "tfr.keras.canned.DNNRankingNetwork.metrics",
+    "tfr.keras.layers.DocumentInteractionAttention.name": "tfr.keras.canned.DNNRankingNetwork.name",
+    "tfr.keras.layers.DocumentInteractionAttention.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
+    "tfr.keras.layers.DocumentInteractionAttention.non_trainable_weights": "tfr.keras.canned.DNNRankingNetwork.non_trainable_weights",
+    "tfr.keras.layers.DocumentInteractionAttention.output": "tfr.keras.canned.DNNRankingNetwork.output",
+    "tfr.keras.layers.DocumentInteractionAttention.set_weights": "tfr.keras.canned.DNNRankingNetwork.set_weights",
+    "tfr.keras.layers.DocumentInteractionAttention.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
+    "tfr.keras.layers.DocumentInteractionAttention.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
+    "tfr.keras.layers.DocumentInteractionAttention.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
+    "tfr.keras.layers.DocumentInteractionAttention.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.layers.DocumentInteractionAttention.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
+    "tfr.keras.layers.DocumentInteractionAttention.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
+    "tfr.keras.layers.FlattenList.__call__": "tfr.keras.canned.DNNRankingNetwork.__call__",
+    "tfr.keras.layers.FlattenList.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
+    "tfr.keras.layers.FlattenList.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
+    "tfr.keras.layers.FlattenList.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
+    "tfr.keras.layers.FlattenList.__le__": "tfr.keras.canned.DNNRankingNetwork.__le__",
+    "tfr.keras.layers.FlattenList.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
+    "tfr.keras.layers.FlattenList.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
+    "tfr.keras.layers.FlattenList.__new__": "tfr.keras.canned.DNNRankingNetwork.__new__",
+    "tfr.keras.layers.FlattenList.activity_regularizer": "tfr.keras.canned.DNNRankingNetwork.activity_regularizer",
+    "tfr.keras.layers.FlattenList.add_loss": "tfr.keras.canned.DNNRankingNetwork.add_loss",
+    "tfr.keras.layers.FlattenList.add_metric": "tfr.keras.canned.DNNRankingNetwork.add_metric",
+    "tfr.keras.layers.FlattenList.add_weight": "tfr.keras.canned.DNNRankingNetwork.add_weight",
+    "tfr.keras.layers.FlattenList.build": "tfr.keras.canned.DNNRankingNetwork.build",
+    "tfr.keras.layers.FlattenList.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
+    "tfr.keras.layers.FlattenList.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
+    "tfr.keras.layers.FlattenList.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
+    "tfr.keras.layers.FlattenList.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
+    "tfr.keras.layers.FlattenList.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
+    "tfr.keras.layers.FlattenList.dtype": "tfr.keras.canned.DNNRankingNetwork.dtype",
+    "tfr.keras.layers.FlattenList.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
+    "tfr.keras.layers.FlattenList.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
+    "tfr.keras.layers.FlattenList.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
+    "tfr.keras.layers.FlattenList.input": "tfr.keras.canned.DNNRankingNetwork.input",
+    "tfr.keras.layers.FlattenList.input_spec": "tfr.keras.canned.DNNRankingNetwork.input_spec",
+    "tfr.keras.layers.FlattenList.losses": "tfr.keras.canned.DNNRankingNetwork.losses",
+    "tfr.keras.layers.FlattenList.metrics": "tfr.keras.canned.DNNRankingNetwork.metrics",
+    "tfr.keras.layers.FlattenList.name": "tfr.keras.canned.DNNRankingNetwork.name",
+    "tfr.keras.layers.FlattenList.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
+    "tfr.keras.layers.FlattenList.non_trainable_weights": "tfr.keras.canned.DNNRankingNetwork.non_trainable_weights",
+    "tfr.keras.layers.FlattenList.output": "tfr.keras.canned.DNNRankingNetwork.output",
+    "tfr.keras.layers.FlattenList.set_weights": "tfr.keras.canned.DNNRankingNetwork.set_weights",
+    "tfr.keras.layers.FlattenList.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
+    "tfr.keras.layers.FlattenList.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
+    "tfr.keras.layers.FlattenList.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
+    "tfr.keras.layers.FlattenList.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.layers.FlattenList.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
+    "tfr.keras.layers.FlattenList.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
+    "tfr.keras.layers.RestoreList.__call__": "tfr.keras.canned.DNNRankingNetwork.__call__",
+    "tfr.keras.layers.RestoreList.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
+    "tfr.keras.layers.RestoreList.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
+    "tfr.keras.layers.RestoreList.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
+    "tfr.keras.layers.RestoreList.__le__": "tfr.keras.canned.DNNRankingNetwork.__le__",
+    "tfr.keras.layers.RestoreList.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
+    "tfr.keras.layers.RestoreList.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
+    "tfr.keras.layers.RestoreList.__new__": "tfr.keras.canned.DNNRankingNetwork.__new__",
+    "tfr.keras.layers.RestoreList.activity_regularizer": "tfr.keras.canned.DNNRankingNetwork.activity_regularizer",
+    "tfr.keras.layers.RestoreList.add_loss": "tfr.keras.canned.DNNRankingNetwork.add_loss",
+    "tfr.keras.layers.RestoreList.add_metric": "tfr.keras.canned.DNNRankingNetwork.add_metric",
+    "tfr.keras.layers.RestoreList.add_weight": "tfr.keras.canned.DNNRankingNetwork.add_weight",
+    "tfr.keras.layers.RestoreList.build": "tfr.keras.canned.DNNRankingNetwork.build",
+    "tfr.keras.layers.RestoreList.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
+    "tfr.keras.layers.RestoreList.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
+    "tfr.keras.layers.RestoreList.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
+    "tfr.keras.layers.RestoreList.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
+    "tfr.keras.layers.RestoreList.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
+    "tfr.keras.layers.RestoreList.dtype": "tfr.keras.canned.DNNRankingNetwork.dtype",
+    "tfr.keras.layers.RestoreList.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
+    "tfr.keras.layers.RestoreList.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
+    "tfr.keras.layers.RestoreList.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
+    "tfr.keras.layers.RestoreList.input": "tfr.keras.canned.DNNRankingNetwork.input",
+    "tfr.keras.layers.RestoreList.input_spec": "tfr.keras.canned.DNNRankingNetwork.input_spec",
+    "tfr.keras.layers.RestoreList.losses": "tfr.keras.canned.DNNRankingNetwork.losses",
+    "tfr.keras.layers.RestoreList.metrics": "tfr.keras.canned.DNNRankingNetwork.metrics",
+    "tfr.keras.layers.RestoreList.name": "tfr.keras.canned.DNNRankingNetwork.name",
+    "tfr.keras.layers.RestoreList.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
+    "tfr.keras.layers.RestoreList.non_trainable_weights": "tfr.keras.canned.DNNRankingNetwork.non_trainable_weights",
+    "tfr.keras.layers.RestoreList.output": "tfr.keras.canned.DNNRankingNetwork.output",
+    "tfr.keras.layers.RestoreList.set_weights": "tfr.keras.canned.DNNRankingNetwork.set_weights",
+    "tfr.keras.layers.RestoreList.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
+    "tfr.keras.layers.RestoreList.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
+    "tfr.keras.layers.RestoreList.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
+    "tfr.keras.layers.RestoreList.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.layers.RestoreList.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
+    "tfr.keras.layers.RestoreList.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.losses.ApproxMRRLoss.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.keras.losses.ApproxMRRLoss.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
     "tfr.keras.losses.ApproxMRRLoss.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
@@ -301,6 +467,22 @@
     "tfr.keras.losses.ApproxNDCGLoss.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
     "tfr.keras.losses.ApproxNDCGLoss.call": "tfr.keras.losses.ApproxMRRLoss.call",
     "tfr.keras.losses.ApproxNDCGLoss.get_config": "tfr.keras.losses.ApproxMRRLoss.get_config",
+    "tfr.keras.losses.ClickEMLoss.__call__": "tfr.keras.losses.ApproxMRRLoss.__call__",
+    "tfr.keras.losses.ClickEMLoss.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
+    "tfr.keras.losses.ClickEMLoss.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
+    "tfr.keras.losses.ClickEMLoss.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
+    "tfr.keras.losses.ClickEMLoss.__le__": "tfr.keras.canned.DNNRankingNetwork.__le__",
+    "tfr.keras.losses.ClickEMLoss.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
+    "tfr.keras.losses.ClickEMLoss.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
+    "tfr.keras.losses.ClickEMLoss.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
+    "tfr.keras.losses.ClickEMLoss.call": "tfr.keras.losses.ApproxMRRLoss.call",
+    "tfr.keras.losses.DCGLambdaWeight.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
+    "tfr.keras.losses.DCGLambdaWeight.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
+    "tfr.keras.losses.DCGLambdaWeight.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
+    "tfr.keras.losses.DCGLambdaWeight.__le__": "tfr.keras.canned.DNNRankingNetwork.__le__",
+    "tfr.keras.losses.DCGLambdaWeight.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
+    "tfr.keras.losses.DCGLambdaWeight.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
+    "tfr.keras.losses.DCGLambdaWeight.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
     "tfr.keras.losses.GumbelApproxNDCGLoss.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.keras.losses.GumbelApproxNDCGLoss.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
     "tfr.keras.losses.GumbelApproxNDCGLoss.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
@@ -309,7 +491,13 @@
     "tfr.keras.losses.GumbelApproxNDCGLoss.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
     "tfr.keras.losses.GumbelApproxNDCGLoss.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
     "tfr.keras.losses.GumbelApproxNDCGLoss.call": "tfr.keras.losses.ApproxMRRLoss.call",
-    "tfr.keras.losses.GumbelApproxNDCGLoss.get_config": "tfr.keras.losses.ApproxMRRLoss.get_config",
+    "tfr.keras.losses.ListMLELambdaWeight.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
+    "tfr.keras.losses.ListMLELambdaWeight.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
+    "tfr.keras.losses.ListMLELambdaWeight.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
+    "tfr.keras.losses.ListMLELambdaWeight.__le__": "tfr.keras.canned.DNNRankingNetwork.__le__",
+    "tfr.keras.losses.ListMLELambdaWeight.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
+    "tfr.keras.losses.ListMLELambdaWeight.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
+    "tfr.keras.losses.ListMLELambdaWeight.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
     "tfr.keras.losses.ListMLELoss.__call__": "tfr.keras.losses.ApproxMRRLoss.__call__",
     "tfr.keras.losses.ListMLELoss.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.keras.losses.ListMLELoss.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
@@ -329,7 +517,16 @@
     "tfr.keras.losses.MeanSquaredLoss.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
     "tfr.keras.losses.MeanSquaredLoss.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
     "tfr.keras.losses.MeanSquaredLoss.call": "tfr.keras.losses.ApproxMRRLoss.call",
-    "tfr.keras.losses.MeanSquaredLoss.get_config": "tfr.keras.losses.ApproxMRRLoss.get_config",
+    "tfr.keras.losses.NDCGLambdaWeight.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
+    "tfr.keras.losses.NDCGLambdaWeight.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
+    "tfr.keras.losses.NDCGLambdaWeight.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
+    "tfr.keras.losses.NDCGLambdaWeight.__le__": "tfr.keras.canned.DNNRankingNetwork.__le__",
+    "tfr.keras.losses.NDCGLambdaWeight.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
+    "tfr.keras.losses.NDCGLambdaWeight.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
+    "tfr.keras.losses.NDCGLambdaWeight.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
+    "tfr.keras.losses.NDCGLambdaWeight.get_config": "tfr.keras.losses.DCGLambdaWeight.get_config",
+    "tfr.keras.losses.NDCGLambdaWeight.individual_weights": "tfr.keras.losses.DCGLambdaWeight.individual_weights",
+    "tfr.keras.losses.NDCGLambdaWeight.pair_weights": "tfr.keras.losses.DCGLambdaWeight.pair_weights",
     "tfr.keras.losses.PairwiseHingeLoss.__call__": "tfr.keras.losses.ApproxMRRLoss.__call__",
     "tfr.keras.losses.PairwiseHingeLoss.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.keras.losses.PairwiseHingeLoss.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
@@ -338,7 +535,6 @@
     "tfr.keras.losses.PairwiseHingeLoss.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
     "tfr.keras.losses.PairwiseHingeLoss.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
     "tfr.keras.losses.PairwiseHingeLoss.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
-    "tfr.keras.losses.PairwiseHingeLoss.get_config": "tfr.keras.losses.ApproxMRRLoss.get_config",
     "tfr.keras.losses.PairwiseLogisticLoss.__call__": "tfr.keras.losses.ApproxMRRLoss.__call__",
     "tfr.keras.losses.PairwiseLogisticLoss.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.keras.losses.PairwiseLogisticLoss.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
@@ -348,7 +544,7 @@
     "tfr.keras.losses.PairwiseLogisticLoss.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
     "tfr.keras.losses.PairwiseLogisticLoss.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
     "tfr.keras.losses.PairwiseLogisticLoss.call": "tfr.keras.losses.PairwiseHingeLoss.call",
-    "tfr.keras.losses.PairwiseLogisticLoss.get_config": "tfr.keras.losses.ApproxMRRLoss.get_config",
+    "tfr.keras.losses.PairwiseLogisticLoss.get_config": "tfr.keras.losses.PairwiseHingeLoss.get_config",
     "tfr.keras.losses.PairwiseSoftZeroOneLoss.__call__": "tfr.keras.losses.ApproxMRRLoss.__call__",
     "tfr.keras.losses.PairwiseSoftZeroOneLoss.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.keras.losses.PairwiseSoftZeroOneLoss.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
@@ -358,7 +554,14 @@
     "tfr.keras.losses.PairwiseSoftZeroOneLoss.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
     "tfr.keras.losses.PairwiseSoftZeroOneLoss.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
     "tfr.keras.losses.PairwiseSoftZeroOneLoss.call": "tfr.keras.losses.PairwiseHingeLoss.call",
-    "tfr.keras.losses.PairwiseSoftZeroOneLoss.get_config": "tfr.keras.losses.ApproxMRRLoss.get_config",
+    "tfr.keras.losses.PairwiseSoftZeroOneLoss.get_config": "tfr.keras.losses.PairwiseHingeLoss.get_config",
+    "tfr.keras.losses.PrecisionLambdaWeight.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
+    "tfr.keras.losses.PrecisionLambdaWeight.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
+    "tfr.keras.losses.PrecisionLambdaWeight.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
+    "tfr.keras.losses.PrecisionLambdaWeight.__le__": "tfr.keras.canned.DNNRankingNetwork.__le__",
+    "tfr.keras.losses.PrecisionLambdaWeight.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
+    "tfr.keras.losses.PrecisionLambdaWeight.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
+    "tfr.keras.losses.PrecisionLambdaWeight.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
     "tfr.keras.losses.RankingLossKey.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.keras.losses.RankingLossKey.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
     "tfr.keras.losses.RankingLossKey.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
@@ -375,7 +578,7 @@
     "tfr.keras.losses.SigmoidCrossEntropyLoss.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
     "tfr.keras.losses.SigmoidCrossEntropyLoss.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
     "tfr.keras.losses.SigmoidCrossEntropyLoss.call": "tfr.keras.losses.ApproxMRRLoss.call",
-    "tfr.keras.losses.SigmoidCrossEntropyLoss.get_config": "tfr.keras.losses.ApproxMRRLoss.get_config",
+    "tfr.keras.losses.SigmoidCrossEntropyLoss.get_config": "tfr.keras.losses.MeanSquaredLoss.get_config",
     "tfr.keras.losses.SoftmaxLoss.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.keras.losses.SoftmaxLoss.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
     "tfr.keras.losses.SoftmaxLoss.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
@@ -405,10 +608,12 @@
     "tfr.keras.metrics.ARPMetric.add_loss": "tfr.keras.canned.DNNRankingNetwork.add_loss",
     "tfr.keras.metrics.ARPMetric.add_metric": "tfr.keras.canned.DNNRankingNetwork.add_metric",
     "tfr.keras.metrics.ARPMetric.build": "tfr.keras.canned.DNNRankingNetwork.build",
+    "tfr.keras.metrics.ARPMetric.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.metrics.ARPMetric.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.metrics.ARPMetric.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
     "tfr.keras.metrics.ARPMetric.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
     "tfr.keras.metrics.ARPMetric.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
+    "tfr.keras.metrics.ARPMetric.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.metrics.ARPMetric.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.metrics.ARPMetric.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
     "tfr.keras.metrics.ARPMetric.input": "tfr.keras.canned.DNNRankingNetwork.input",
@@ -417,13 +622,12 @@
     "tfr.keras.metrics.ARPMetric.metrics": "tfr.keras.canned.DNNRankingNetwork.metrics",
     "tfr.keras.metrics.ARPMetric.name": "tfr.keras.canned.DNNRankingNetwork.name",
     "tfr.keras.metrics.ARPMetric.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
-    "tfr.keras.metrics.ARPMetric.non_trainable_weights": "tfr.keras.canned.DNNRankingNetwork.non_trainable_weights",
     "tfr.keras.metrics.ARPMetric.output": "tfr.keras.canned.DNNRankingNetwork.output",
     "tfr.keras.metrics.ARPMetric.set_weights": "tfr.keras.canned.DNNRankingNetwork.set_weights",
     "tfr.keras.metrics.ARPMetric.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
     "tfr.keras.metrics.ARPMetric.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
     "tfr.keras.metrics.ARPMetric.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
-    "tfr.keras.metrics.ARPMetric.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.metrics.ARPMetric.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.metrics.ARPMetric.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.metrics.AlphaDCGMetric.__call__": "tfr.keras.metrics.ARPMetric.__call__",
     "tfr.keras.metrics.AlphaDCGMetric.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
@@ -439,11 +643,13 @@
     "tfr.keras.metrics.AlphaDCGMetric.add_weight": "tfr.keras.metrics.ARPMetric.add_weight",
     "tfr.keras.metrics.AlphaDCGMetric.build": "tfr.keras.canned.DNNRankingNetwork.build",
     "tfr.keras.metrics.AlphaDCGMetric.call": "tfr.keras.metrics.ARPMetric.call",
+    "tfr.keras.metrics.AlphaDCGMetric.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.metrics.AlphaDCGMetric.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.metrics.AlphaDCGMetric.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
     "tfr.keras.metrics.AlphaDCGMetric.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
     "tfr.keras.metrics.AlphaDCGMetric.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.metrics.AlphaDCGMetric.dtype": "tfr.keras.metrics.ARPMetric.dtype",
+    "tfr.keras.metrics.AlphaDCGMetric.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.metrics.AlphaDCGMetric.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.metrics.AlphaDCGMetric.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
     "tfr.keras.metrics.AlphaDCGMetric.input": "tfr.keras.canned.DNNRankingNetwork.input",
@@ -452,7 +658,7 @@
     "tfr.keras.metrics.AlphaDCGMetric.metrics": "tfr.keras.canned.DNNRankingNetwork.metrics",
     "tfr.keras.metrics.AlphaDCGMetric.name": "tfr.keras.canned.DNNRankingNetwork.name",
     "tfr.keras.metrics.AlphaDCGMetric.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
-    "tfr.keras.metrics.AlphaDCGMetric.non_trainable_weights": "tfr.keras.canned.DNNRankingNetwork.non_trainable_weights",
+    "tfr.keras.metrics.AlphaDCGMetric.non_trainable_weights": "tfr.keras.metrics.ARPMetric.non_trainable_weights",
     "tfr.keras.metrics.AlphaDCGMetric.output": "tfr.keras.canned.DNNRankingNetwork.output",
     "tfr.keras.metrics.AlphaDCGMetric.reset_states": "tfr.keras.metrics.ARPMetric.reset_states",
     "tfr.keras.metrics.AlphaDCGMetric.result": "tfr.keras.metrics.ARPMetric.result",
@@ -460,8 +666,9 @@
     "tfr.keras.metrics.AlphaDCGMetric.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
     "tfr.keras.metrics.AlphaDCGMetric.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
     "tfr.keras.metrics.AlphaDCGMetric.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
-    "tfr.keras.metrics.AlphaDCGMetric.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.metrics.AlphaDCGMetric.trainable_weights": "tfr.keras.metrics.ARPMetric.trainable_weights",
     "tfr.keras.metrics.AlphaDCGMetric.update_state": "tfr.keras.metrics.ARPMetric.update_state",
+    "tfr.keras.metrics.AlphaDCGMetric.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.metrics.AlphaDCGMetric.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.metrics.DCGMetric.__call__": "tfr.keras.metrics.ARPMetric.__call__",
     "tfr.keras.metrics.DCGMetric.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
@@ -477,11 +684,13 @@
     "tfr.keras.metrics.DCGMetric.add_weight": "tfr.keras.metrics.ARPMetric.add_weight",
     "tfr.keras.metrics.DCGMetric.build": "tfr.keras.canned.DNNRankingNetwork.build",
     "tfr.keras.metrics.DCGMetric.call": "tfr.keras.metrics.ARPMetric.call",
+    "tfr.keras.metrics.DCGMetric.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.metrics.DCGMetric.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.metrics.DCGMetric.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
     "tfr.keras.metrics.DCGMetric.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
     "tfr.keras.metrics.DCGMetric.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.metrics.DCGMetric.dtype": "tfr.keras.metrics.ARPMetric.dtype",
+    "tfr.keras.metrics.DCGMetric.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.metrics.DCGMetric.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.metrics.DCGMetric.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
     "tfr.keras.metrics.DCGMetric.input": "tfr.keras.canned.DNNRankingNetwork.input",
@@ -490,7 +699,7 @@
     "tfr.keras.metrics.DCGMetric.metrics": "tfr.keras.canned.DNNRankingNetwork.metrics",
     "tfr.keras.metrics.DCGMetric.name": "tfr.keras.canned.DNNRankingNetwork.name",
     "tfr.keras.metrics.DCGMetric.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
-    "tfr.keras.metrics.DCGMetric.non_trainable_weights": "tfr.keras.canned.DNNRankingNetwork.non_trainable_weights",
+    "tfr.keras.metrics.DCGMetric.non_trainable_weights": "tfr.keras.metrics.ARPMetric.non_trainable_weights",
     "tfr.keras.metrics.DCGMetric.output": "tfr.keras.canned.DNNRankingNetwork.output",
     "tfr.keras.metrics.DCGMetric.reset_states": "tfr.keras.metrics.ARPMetric.reset_states",
     "tfr.keras.metrics.DCGMetric.result": "tfr.keras.metrics.ARPMetric.result",
@@ -498,8 +707,9 @@
     "tfr.keras.metrics.DCGMetric.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
     "tfr.keras.metrics.DCGMetric.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
     "tfr.keras.metrics.DCGMetric.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
-    "tfr.keras.metrics.DCGMetric.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.metrics.DCGMetric.trainable_weights": "tfr.keras.metrics.ARPMetric.trainable_weights",
     "tfr.keras.metrics.DCGMetric.update_state": "tfr.keras.metrics.ARPMetric.update_state",
+    "tfr.keras.metrics.DCGMetric.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.metrics.DCGMetric.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.metrics.MRRMetric.__call__": "tfr.keras.metrics.ARPMetric.__call__",
     "tfr.keras.metrics.MRRMetric.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
@@ -515,11 +725,13 @@
     "tfr.keras.metrics.MRRMetric.add_weight": "tfr.keras.metrics.ARPMetric.add_weight",
     "tfr.keras.metrics.MRRMetric.build": "tfr.keras.canned.DNNRankingNetwork.build",
     "tfr.keras.metrics.MRRMetric.call": "tfr.keras.metrics.ARPMetric.call",
+    "tfr.keras.metrics.MRRMetric.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.metrics.MRRMetric.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.metrics.MRRMetric.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
     "tfr.keras.metrics.MRRMetric.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
     "tfr.keras.metrics.MRRMetric.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.metrics.MRRMetric.dtype": "tfr.keras.metrics.ARPMetric.dtype",
+    "tfr.keras.metrics.MRRMetric.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.metrics.MRRMetric.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.metrics.MRRMetric.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
     "tfr.keras.metrics.MRRMetric.input": "tfr.keras.canned.DNNRankingNetwork.input",
@@ -528,7 +740,7 @@
     "tfr.keras.metrics.MRRMetric.metrics": "tfr.keras.canned.DNNRankingNetwork.metrics",
     "tfr.keras.metrics.MRRMetric.name": "tfr.keras.canned.DNNRankingNetwork.name",
     "tfr.keras.metrics.MRRMetric.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
-    "tfr.keras.metrics.MRRMetric.non_trainable_weights": "tfr.keras.canned.DNNRankingNetwork.non_trainable_weights",
+    "tfr.keras.metrics.MRRMetric.non_trainable_weights": "tfr.keras.metrics.ARPMetric.non_trainable_weights",
     "tfr.keras.metrics.MRRMetric.output": "tfr.keras.canned.DNNRankingNetwork.output",
     "tfr.keras.metrics.MRRMetric.reset_states": "tfr.keras.metrics.ARPMetric.reset_states",
     "tfr.keras.metrics.MRRMetric.result": "tfr.keras.metrics.ARPMetric.result",
@@ -536,8 +748,9 @@
     "tfr.keras.metrics.MRRMetric.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
     "tfr.keras.metrics.MRRMetric.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
     "tfr.keras.metrics.MRRMetric.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
-    "tfr.keras.metrics.MRRMetric.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.metrics.MRRMetric.trainable_weights": "tfr.keras.metrics.ARPMetric.trainable_weights",
     "tfr.keras.metrics.MRRMetric.update_state": "tfr.keras.metrics.ARPMetric.update_state",
+    "tfr.keras.metrics.MRRMetric.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.metrics.MRRMetric.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.__call__": "tfr.keras.metrics.ARPMetric.__call__",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
@@ -553,11 +766,13 @@
     "tfr.keras.metrics.MeanAveragePrecisionMetric.add_weight": "tfr.keras.metrics.ARPMetric.add_weight",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.build": "tfr.keras.canned.DNNRankingNetwork.build",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.call": "tfr.keras.metrics.ARPMetric.call",
+    "tfr.keras.metrics.MeanAveragePrecisionMetric.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.dtype": "tfr.keras.metrics.ARPMetric.dtype",
+    "tfr.keras.metrics.MeanAveragePrecisionMetric.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.input": "tfr.keras.canned.DNNRankingNetwork.input",
@@ -566,7 +781,7 @@
     "tfr.keras.metrics.MeanAveragePrecisionMetric.metrics": "tfr.keras.canned.DNNRankingNetwork.metrics",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.name": "tfr.keras.canned.DNNRankingNetwork.name",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
-    "tfr.keras.metrics.MeanAveragePrecisionMetric.non_trainable_weights": "tfr.keras.canned.DNNRankingNetwork.non_trainable_weights",
+    "tfr.keras.metrics.MeanAveragePrecisionMetric.non_trainable_weights": "tfr.keras.metrics.ARPMetric.non_trainable_weights",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.output": "tfr.keras.canned.DNNRankingNetwork.output",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.reset_states": "tfr.keras.metrics.ARPMetric.reset_states",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.result": "tfr.keras.metrics.ARPMetric.result",
@@ -574,8 +789,9 @@
     "tfr.keras.metrics.MeanAveragePrecisionMetric.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
-    "tfr.keras.metrics.MeanAveragePrecisionMetric.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.metrics.MeanAveragePrecisionMetric.trainable_weights": "tfr.keras.metrics.ARPMetric.trainable_weights",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.update_state": "tfr.keras.metrics.ARPMetric.update_state",
+    "tfr.keras.metrics.MeanAveragePrecisionMetric.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.metrics.NDCGMetric.__call__": "tfr.keras.metrics.ARPMetric.__call__",
     "tfr.keras.metrics.NDCGMetric.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
@@ -591,11 +807,13 @@
     "tfr.keras.metrics.NDCGMetric.add_weight": "tfr.keras.metrics.ARPMetric.add_weight",
     "tfr.keras.metrics.NDCGMetric.build": "tfr.keras.canned.DNNRankingNetwork.build",
     "tfr.keras.metrics.NDCGMetric.call": "tfr.keras.metrics.ARPMetric.call",
+    "tfr.keras.metrics.NDCGMetric.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.metrics.NDCGMetric.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.metrics.NDCGMetric.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
     "tfr.keras.metrics.NDCGMetric.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
     "tfr.keras.metrics.NDCGMetric.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.metrics.NDCGMetric.dtype": "tfr.keras.metrics.ARPMetric.dtype",
+    "tfr.keras.metrics.NDCGMetric.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.metrics.NDCGMetric.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.metrics.NDCGMetric.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
     "tfr.keras.metrics.NDCGMetric.input": "tfr.keras.canned.DNNRankingNetwork.input",
@@ -604,7 +822,7 @@
     "tfr.keras.metrics.NDCGMetric.metrics": "tfr.keras.canned.DNNRankingNetwork.metrics",
     "tfr.keras.metrics.NDCGMetric.name": "tfr.keras.canned.DNNRankingNetwork.name",
     "tfr.keras.metrics.NDCGMetric.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
-    "tfr.keras.metrics.NDCGMetric.non_trainable_weights": "tfr.keras.canned.DNNRankingNetwork.non_trainable_weights",
+    "tfr.keras.metrics.NDCGMetric.non_trainable_weights": "tfr.keras.metrics.ARPMetric.non_trainable_weights",
     "tfr.keras.metrics.NDCGMetric.output": "tfr.keras.canned.DNNRankingNetwork.output",
     "tfr.keras.metrics.NDCGMetric.reset_states": "tfr.keras.metrics.ARPMetric.reset_states",
     "tfr.keras.metrics.NDCGMetric.result": "tfr.keras.metrics.ARPMetric.result",
@@ -612,8 +830,9 @@
     "tfr.keras.metrics.NDCGMetric.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
     "tfr.keras.metrics.NDCGMetric.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
     "tfr.keras.metrics.NDCGMetric.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
-    "tfr.keras.metrics.NDCGMetric.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.metrics.NDCGMetric.trainable_weights": "tfr.keras.metrics.ARPMetric.trainable_weights",
     "tfr.keras.metrics.NDCGMetric.update_state": "tfr.keras.metrics.ARPMetric.update_state",
+    "tfr.keras.metrics.NDCGMetric.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.metrics.NDCGMetric.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.metrics.OPAMetric.__call__": "tfr.keras.metrics.ARPMetric.__call__",
     "tfr.keras.metrics.OPAMetric.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
@@ -629,11 +848,13 @@
     "tfr.keras.metrics.OPAMetric.add_weight": "tfr.keras.metrics.ARPMetric.add_weight",
     "tfr.keras.metrics.OPAMetric.build": "tfr.keras.canned.DNNRankingNetwork.build",
     "tfr.keras.metrics.OPAMetric.call": "tfr.keras.metrics.ARPMetric.call",
+    "tfr.keras.metrics.OPAMetric.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.metrics.OPAMetric.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.metrics.OPAMetric.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
     "tfr.keras.metrics.OPAMetric.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
     "tfr.keras.metrics.OPAMetric.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.metrics.OPAMetric.dtype": "tfr.keras.metrics.ARPMetric.dtype",
+    "tfr.keras.metrics.OPAMetric.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.metrics.OPAMetric.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.metrics.OPAMetric.get_config": "tfr.keras.metrics.ARPMetric.get_config",
     "tfr.keras.metrics.OPAMetric.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
@@ -643,7 +864,7 @@
     "tfr.keras.metrics.OPAMetric.metrics": "tfr.keras.canned.DNNRankingNetwork.metrics",
     "tfr.keras.metrics.OPAMetric.name": "tfr.keras.canned.DNNRankingNetwork.name",
     "tfr.keras.metrics.OPAMetric.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
-    "tfr.keras.metrics.OPAMetric.non_trainable_weights": "tfr.keras.canned.DNNRankingNetwork.non_trainable_weights",
+    "tfr.keras.metrics.OPAMetric.non_trainable_weights": "tfr.keras.metrics.ARPMetric.non_trainable_weights",
     "tfr.keras.metrics.OPAMetric.output": "tfr.keras.canned.DNNRankingNetwork.output",
     "tfr.keras.metrics.OPAMetric.reset_states": "tfr.keras.metrics.ARPMetric.reset_states",
     "tfr.keras.metrics.OPAMetric.result": "tfr.keras.metrics.ARPMetric.result",
@@ -651,8 +872,9 @@
     "tfr.keras.metrics.OPAMetric.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
     "tfr.keras.metrics.OPAMetric.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
     "tfr.keras.metrics.OPAMetric.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
-    "tfr.keras.metrics.OPAMetric.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.metrics.OPAMetric.trainable_weights": "tfr.keras.metrics.ARPMetric.trainable_weights",
     "tfr.keras.metrics.OPAMetric.update_state": "tfr.keras.metrics.ARPMetric.update_state",
+    "tfr.keras.metrics.OPAMetric.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.metrics.OPAMetric.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.metrics.PrecisionIAMetric.__call__": "tfr.keras.metrics.ARPMetric.__call__",
     "tfr.keras.metrics.PrecisionIAMetric.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
@@ -668,11 +890,13 @@
     "tfr.keras.metrics.PrecisionIAMetric.add_weight": "tfr.keras.metrics.ARPMetric.add_weight",
     "tfr.keras.metrics.PrecisionIAMetric.build": "tfr.keras.canned.DNNRankingNetwork.build",
     "tfr.keras.metrics.PrecisionIAMetric.call": "tfr.keras.metrics.ARPMetric.call",
+    "tfr.keras.metrics.PrecisionIAMetric.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.metrics.PrecisionIAMetric.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.metrics.PrecisionIAMetric.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
     "tfr.keras.metrics.PrecisionIAMetric.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
     "tfr.keras.metrics.PrecisionIAMetric.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.metrics.PrecisionIAMetric.dtype": "tfr.keras.metrics.ARPMetric.dtype",
+    "tfr.keras.metrics.PrecisionIAMetric.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.metrics.PrecisionIAMetric.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.metrics.PrecisionIAMetric.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
     "tfr.keras.metrics.PrecisionIAMetric.input": "tfr.keras.canned.DNNRankingNetwork.input",
@@ -681,7 +905,7 @@
     "tfr.keras.metrics.PrecisionIAMetric.metrics": "tfr.keras.canned.DNNRankingNetwork.metrics",
     "tfr.keras.metrics.PrecisionIAMetric.name": "tfr.keras.canned.DNNRankingNetwork.name",
     "tfr.keras.metrics.PrecisionIAMetric.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
-    "tfr.keras.metrics.PrecisionIAMetric.non_trainable_weights": "tfr.keras.canned.DNNRankingNetwork.non_trainable_weights",
+    "tfr.keras.metrics.PrecisionIAMetric.non_trainable_weights": "tfr.keras.metrics.ARPMetric.non_trainable_weights",
     "tfr.keras.metrics.PrecisionIAMetric.output": "tfr.keras.canned.DNNRankingNetwork.output",
     "tfr.keras.metrics.PrecisionIAMetric.reset_states": "tfr.keras.metrics.ARPMetric.reset_states",
     "tfr.keras.metrics.PrecisionIAMetric.result": "tfr.keras.metrics.ARPMetric.result",
@@ -689,8 +913,9 @@
     "tfr.keras.metrics.PrecisionIAMetric.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
     "tfr.keras.metrics.PrecisionIAMetric.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
     "tfr.keras.metrics.PrecisionIAMetric.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
-    "tfr.keras.metrics.PrecisionIAMetric.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.metrics.PrecisionIAMetric.trainable_weights": "tfr.keras.metrics.ARPMetric.trainable_weights",
     "tfr.keras.metrics.PrecisionIAMetric.update_state": "tfr.keras.metrics.ARPMetric.update_state",
+    "tfr.keras.metrics.PrecisionIAMetric.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.metrics.PrecisionIAMetric.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.metrics.PrecisionMetric.__call__": "tfr.keras.metrics.ARPMetric.__call__",
     "tfr.keras.metrics.PrecisionMetric.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
@@ -706,11 +931,13 @@
     "tfr.keras.metrics.PrecisionMetric.add_weight": "tfr.keras.metrics.ARPMetric.add_weight",
     "tfr.keras.metrics.PrecisionMetric.build": "tfr.keras.canned.DNNRankingNetwork.build",
     "tfr.keras.metrics.PrecisionMetric.call": "tfr.keras.metrics.ARPMetric.call",
+    "tfr.keras.metrics.PrecisionMetric.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.metrics.PrecisionMetric.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.metrics.PrecisionMetric.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
     "tfr.keras.metrics.PrecisionMetric.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
     "tfr.keras.metrics.PrecisionMetric.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.metrics.PrecisionMetric.dtype": "tfr.keras.metrics.ARPMetric.dtype",
+    "tfr.keras.metrics.PrecisionMetric.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.metrics.PrecisionMetric.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.metrics.PrecisionMetric.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
     "tfr.keras.metrics.PrecisionMetric.input": "tfr.keras.canned.DNNRankingNetwork.input",
@@ -719,7 +946,7 @@
     "tfr.keras.metrics.PrecisionMetric.metrics": "tfr.keras.canned.DNNRankingNetwork.metrics",
     "tfr.keras.metrics.PrecisionMetric.name": "tfr.keras.canned.DNNRankingNetwork.name",
     "tfr.keras.metrics.PrecisionMetric.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
-    "tfr.keras.metrics.PrecisionMetric.non_trainable_weights": "tfr.keras.canned.DNNRankingNetwork.non_trainable_weights",
+    "tfr.keras.metrics.PrecisionMetric.non_trainable_weights": "tfr.keras.metrics.ARPMetric.non_trainable_weights",
     "tfr.keras.metrics.PrecisionMetric.output": "tfr.keras.canned.DNNRankingNetwork.output",
     "tfr.keras.metrics.PrecisionMetric.reset_states": "tfr.keras.metrics.ARPMetric.reset_states",
     "tfr.keras.metrics.PrecisionMetric.result": "tfr.keras.metrics.ARPMetric.result",
@@ -727,12 +954,59 @@
     "tfr.keras.metrics.PrecisionMetric.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
     "tfr.keras.metrics.PrecisionMetric.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
     "tfr.keras.metrics.PrecisionMetric.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
-    "tfr.keras.metrics.PrecisionMetric.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.metrics.PrecisionMetric.trainable_weights": "tfr.keras.metrics.ARPMetric.trainable_weights",
     "tfr.keras.metrics.PrecisionMetric.update_state": "tfr.keras.metrics.ARPMetric.update_state",
+    "tfr.keras.metrics.PrecisionMetric.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.metrics.PrecisionMetric.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
-    "tfr.keras.metrics.absolute_import": "tfr.keras.losses.absolute_import",
-    "tfr.keras.metrics.division": "tfr.keras.losses.division",
-    "tfr.keras.metrics.print_function": "tfr.keras.losses.print_function",
+    "tfr.keras.metrics.RankingMetricKey.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
+    "tfr.keras.metrics.RankingMetricKey.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
+    "tfr.keras.metrics.RankingMetricKey.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
+    "tfr.keras.metrics.RankingMetricKey.__init__": "tfr.keras.losses.RankingLossKey.__init__",
+    "tfr.keras.metrics.RankingMetricKey.__le__": "tfr.keras.canned.DNNRankingNetwork.__le__",
+    "tfr.keras.metrics.RankingMetricKey.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
+    "tfr.keras.metrics.RankingMetricKey.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
+    "tfr.keras.metrics.RankingMetricKey.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
+    "tfr.keras.metrics.RecallMetric.__call__": "tfr.keras.metrics.ARPMetric.__call__",
+    "tfr.keras.metrics.RecallMetric.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
+    "tfr.keras.metrics.RecallMetric.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
+    "tfr.keras.metrics.RecallMetric.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
+    "tfr.keras.metrics.RecallMetric.__le__": "tfr.keras.canned.DNNRankingNetwork.__le__",
+    "tfr.keras.metrics.RecallMetric.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
+    "tfr.keras.metrics.RecallMetric.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
+    "tfr.keras.metrics.RecallMetric.__new__": "tfr.keras.metrics.ARPMetric.__new__",
+    "tfr.keras.metrics.RecallMetric.activity_regularizer": "tfr.keras.canned.DNNRankingNetwork.activity_regularizer",
+    "tfr.keras.metrics.RecallMetric.add_loss": "tfr.keras.canned.DNNRankingNetwork.add_loss",
+    "tfr.keras.metrics.RecallMetric.add_metric": "tfr.keras.canned.DNNRankingNetwork.add_metric",
+    "tfr.keras.metrics.RecallMetric.add_weight": "tfr.keras.metrics.ARPMetric.add_weight",
+    "tfr.keras.metrics.RecallMetric.build": "tfr.keras.canned.DNNRankingNetwork.build",
+    "tfr.keras.metrics.RecallMetric.call": "tfr.keras.metrics.ARPMetric.call",
+    "tfr.keras.metrics.RecallMetric.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
+    "tfr.keras.metrics.RecallMetric.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
+    "tfr.keras.metrics.RecallMetric.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
+    "tfr.keras.metrics.RecallMetric.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
+    "tfr.keras.metrics.RecallMetric.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
+    "tfr.keras.metrics.RecallMetric.dtype": "tfr.keras.metrics.ARPMetric.dtype",
+    "tfr.keras.metrics.RecallMetric.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
+    "tfr.keras.metrics.RecallMetric.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
+    "tfr.keras.metrics.RecallMetric.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
+    "tfr.keras.metrics.RecallMetric.input": "tfr.keras.canned.DNNRankingNetwork.input",
+    "tfr.keras.metrics.RecallMetric.input_spec": "tfr.keras.canned.DNNRankingNetwork.input_spec",
+    "tfr.keras.metrics.RecallMetric.losses": "tfr.keras.canned.DNNRankingNetwork.losses",
+    "tfr.keras.metrics.RecallMetric.metrics": "tfr.keras.canned.DNNRankingNetwork.metrics",
+    "tfr.keras.metrics.RecallMetric.name": "tfr.keras.canned.DNNRankingNetwork.name",
+    "tfr.keras.metrics.RecallMetric.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
+    "tfr.keras.metrics.RecallMetric.non_trainable_weights": "tfr.keras.metrics.ARPMetric.non_trainable_weights",
+    "tfr.keras.metrics.RecallMetric.output": "tfr.keras.canned.DNNRankingNetwork.output",
+    "tfr.keras.metrics.RecallMetric.reset_states": "tfr.keras.metrics.ARPMetric.reset_states",
+    "tfr.keras.metrics.RecallMetric.result": "tfr.keras.metrics.ARPMetric.result",
+    "tfr.keras.metrics.RecallMetric.set_weights": "tfr.keras.canned.DNNRankingNetwork.set_weights",
+    "tfr.keras.metrics.RecallMetric.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
+    "tfr.keras.metrics.RecallMetric.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
+    "tfr.keras.metrics.RecallMetric.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
+    "tfr.keras.metrics.RecallMetric.trainable_weights": "tfr.keras.metrics.ARPMetric.trainable_weights",
+    "tfr.keras.metrics.RecallMetric.update_state": "tfr.keras.metrics.ARPMetric.update_state",
+    "tfr.keras.metrics.RecallMetric.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
+    "tfr.keras.metrics.RecallMetric.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.network.RankingNetwork.__call__": "tfr.keras.canned.DNNRankingNetwork.__call__",
     "tfr.keras.network.RankingNetwork.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.keras.network.RankingNetwork.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
@@ -746,11 +1020,13 @@
     "tfr.keras.network.RankingNetwork.add_metric": "tfr.keras.canned.DNNRankingNetwork.add_metric",
     "tfr.keras.network.RankingNetwork.add_weight": "tfr.keras.canned.DNNRankingNetwork.add_weight",
     "tfr.keras.network.RankingNetwork.build": "tfr.keras.canned.DNNRankingNetwork.build",
+    "tfr.keras.network.RankingNetwork.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.network.RankingNetwork.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.network.RankingNetwork.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
     "tfr.keras.network.RankingNetwork.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
     "tfr.keras.network.RankingNetwork.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.network.RankingNetwork.dtype": "tfr.keras.canned.DNNRankingNetwork.dtype",
+    "tfr.keras.network.RankingNetwork.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.network.RankingNetwork.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.network.RankingNetwork.get_weights": "tfr.keras.canned.DNNRankingNetwork.get_weights",
     "tfr.keras.network.RankingNetwork.input": "tfr.keras.canned.DNNRankingNetwork.input",
@@ -766,6 +1042,7 @@
     "tfr.keras.network.RankingNetwork.supports_masking": "tfr.keras.canned.DNNRankingNetwork.supports_masking",
     "tfr.keras.network.RankingNetwork.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
     "tfr.keras.network.RankingNetwork.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
+    "tfr.keras.network.RankingNetwork.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.network.RankingNetwork.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
     "tfr.keras.network.UnivariateRankingNetwork.__call__": "tfr.keras.canned.DNNRankingNetwork.__call__",
     "tfr.keras.network.UnivariateRankingNetwork.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
@@ -781,12 +1058,14 @@
     "tfr.keras.network.UnivariateRankingNetwork.add_weight": "tfr.keras.canned.DNNRankingNetwork.add_weight",
     "tfr.keras.network.UnivariateRankingNetwork.build": "tfr.keras.canned.DNNRankingNetwork.build",
     "tfr.keras.network.UnivariateRankingNetwork.call": "tfr.keras.network.RankingNetwork.call",
+    "tfr.keras.network.UnivariateRankingNetwork.compute_dtype": "tfr.keras.canned.DNNRankingNetwork.compute_dtype",
     "tfr.keras.network.UnivariateRankingNetwork.compute_mask": "tfr.keras.canned.DNNRankingNetwork.compute_mask",
     "tfr.keras.network.UnivariateRankingNetwork.compute_output_shape": "tfr.keras.canned.DNNRankingNetwork.compute_output_shape",
     "tfr.keras.network.UnivariateRankingNetwork.compute_output_signature": "tfr.keras.canned.DNNRankingNetwork.compute_output_signature",
     "tfr.keras.network.UnivariateRankingNetwork.context_feature_columns": "tfr.keras.network.RankingNetwork.context_feature_columns",
     "tfr.keras.network.UnivariateRankingNetwork.count_params": "tfr.keras.canned.DNNRankingNetwork.count_params",
     "tfr.keras.network.UnivariateRankingNetwork.dtype": "tfr.keras.canned.DNNRankingNetwork.dtype",
+    "tfr.keras.network.UnivariateRankingNetwork.dtype_policy": "tfr.keras.canned.DNNRankingNetwork.dtype_policy",
     "tfr.keras.network.UnivariateRankingNetwork.dynamic": "tfr.keras.canned.DNNRankingNetwork.dynamic",
     "tfr.keras.network.UnivariateRankingNetwork.example_feature_columns": "tfr.keras.network.RankingNetwork.example_feature_columns",
     "tfr.keras.network.UnivariateRankingNetwork.get_config": "tfr.keras.network.RankingNetwork.get_config",
@@ -805,7 +1084,17 @@
     "tfr.keras.network.UnivariateRankingNetwork.trainable": "tfr.keras.canned.DNNRankingNetwork.trainable",
     "tfr.keras.network.UnivariateRankingNetwork.trainable_weights": "tfr.keras.canned.DNNRankingNetwork.trainable_weights",
     "tfr.keras.network.UnivariateRankingNetwork.transform": "tfr.keras.network.RankingNetwork.transform",
+    "tfr.keras.network.UnivariateRankingNetwork.variable_dtype": "tfr.keras.canned.DNNRankingNetwork.variable_dtype",
     "tfr.keras.network.UnivariateRankingNetwork.weights": "tfr.keras.canned.DNNRankingNetwork.weights",
+    "tfr.keras.saved_model.Signatures.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
+    "tfr.keras.saved_model.Signatures.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
+    "tfr.keras.saved_model.Signatures.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
+    "tfr.keras.saved_model.Signatures.__le__": "tfr.keras.canned.DNNRankingNetwork.__le__",
+    "tfr.keras.saved_model.Signatures.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
+    "tfr.keras.saved_model.Signatures.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
+    "tfr.keras.saved_model.Signatures.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
+    "tfr.keras.saved_model.Signatures.name_scope": "tfr.keras.canned.DNNRankingNetwork.name_scope",
+    "tfr.keras.saved_model.Signatures.submodules": "tfr.keras.canned.DNNRankingNetwork.submodules",
     "tfr.losses.RankingLossKey.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.losses.RankingLossKey.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
     "tfr.losses.RankingLossKey.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
@@ -814,9 +1103,9 @@
     "tfr.losses.RankingLossKey.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
     "tfr.losses.RankingLossKey.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
     "tfr.losses.RankingLossKey.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
-    "tfr.losses.absolute_import": "tfr.keras.losses.absolute_import",
-    "tfr.losses.division": "tfr.keras.losses.division",
-    "tfr.losses.print_function": "tfr.keras.losses.print_function",
+    "tfr.losses.absolute_import": "tfr.keras.metrics.absolute_import",
+    "tfr.losses.division": "tfr.keras.metrics.division",
+    "tfr.losses.print_function": "tfr.keras.metrics.print_function",
     "tfr.metrics.RankingMetricKey.__eq__": "tfr.keras.canned.DNNRankingNetwork.__eq__",
     "tfr.metrics.RankingMetricKey.__ge__": "tfr.keras.canned.DNNRankingNetwork.__ge__",
     "tfr.metrics.RankingMetricKey.__gt__": "tfr.keras.canned.DNNRankingNetwork.__gt__",
@@ -825,15 +1114,15 @@
     "tfr.metrics.RankingMetricKey.__lt__": "tfr.keras.canned.DNNRankingNetwork.__lt__",
     "tfr.metrics.RankingMetricKey.__ne__": "tfr.keras.canned.DNNRankingNetwork.__ne__",
     "tfr.metrics.RankingMetricKey.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
-    "tfr.metrics.absolute_import": "tfr.keras.losses.absolute_import",
-    "tfr.metrics.division": "tfr.keras.losses.division",
-    "tfr.metrics.print_function": "tfr.keras.losses.print_function",
-    "tfr.model.absolute_import": "tfr.keras.losses.absolute_import",
-    "tfr.model.division": "tfr.keras.losses.division",
-    "tfr.model.print_function": "tfr.keras.losses.print_function",
-    "tfr.utils.absolute_import": "tfr.keras.losses.absolute_import",
-    "tfr.utils.division": "tfr.keras.losses.division",
-    "tfr.utils.print_function": "tfr.keras.losses.print_function"
+    "tfr.metrics.absolute_import": "tfr.keras.metrics.absolute_import",
+    "tfr.metrics.division": "tfr.keras.metrics.division",
+    "tfr.metrics.print_function": "tfr.keras.metrics.print_function",
+    "tfr.model.absolute_import": "tfr.keras.metrics.absolute_import",
+    "tfr.model.division": "tfr.keras.metrics.division",
+    "tfr.model.print_function": "tfr.keras.metrics.print_function",
+    "tfr.utils.absolute_import": "tfr.keras.metrics.absolute_import",
+    "tfr.utils.division": "tfr.keras.metrics.division",
+    "tfr.utils.print_function": "tfr.keras.metrics.print_function"
   },
   "is_fragment": {
     "tfr": false,
@@ -848,11 +1137,13 @@
     "tfr.data.build_ranking_serving_input_receiver_fn": false,
     "tfr.data.build_ranking_serving_input_receiver_fn_with_parsing_fn": false,
     "tfr.data.build_sequence_example_serving_input_receiver_fn": false,
+    "tfr.data.build_tf_example_serving_input_receiver_fn": false,
     "tfr.data.division": true,
     "tfr.data.make_parsing_fn": false,
     "tfr.data.parse_from_example_in_example": false,
     "tfr.data.parse_from_example_list": false,
     "tfr.data.parse_from_sequence_example": false,
+    "tfr.data.parse_from_tf_example": false,
     "tfr.data.print_function": true,
     "tfr.data.read_batched_sequence_example_dataset": false,
     "tfr.estimator": false,
@@ -913,6 +1204,7 @@
     "tfr.ext.tfrbert.TFRBertRankingNetwork.add_weight": true,
     "tfr.ext.tfrbert.TFRBertRankingNetwork.build": true,
     "tfr.ext.tfrbert.TFRBertRankingNetwork.call": true,
+    "tfr.ext.tfrbert.TFRBertRankingNetwork.compute_dtype": true,
     "tfr.ext.tfrbert.TFRBertRankingNetwork.compute_logits": true,
     "tfr.ext.tfrbert.TFRBertRankingNetwork.compute_mask": true,
     "tfr.ext.tfrbert.TFRBertRankingNetwork.compute_output_shape": true,
@@ -920,6 +1212,7 @@
     "tfr.ext.tfrbert.TFRBertRankingNetwork.context_feature_columns": true,
     "tfr.ext.tfrbert.TFRBertRankingNetwork.count_params": true,
     "tfr.ext.tfrbert.TFRBertRankingNetwork.dtype": true,
+    "tfr.ext.tfrbert.TFRBertRankingNetwork.dtype_policy": true,
     "tfr.ext.tfrbert.TFRBertRankingNetwork.dynamic": true,
     "tfr.ext.tfrbert.TFRBertRankingNetwork.example_feature_columns": true,
     "tfr.ext.tfrbert.TFRBertRankingNetwork.from_config": true,
@@ -940,6 +1233,7 @@
     "tfr.ext.tfrbert.TFRBertRankingNetwork.trainable": true,
     "tfr.ext.tfrbert.TFRBertRankingNetwork.trainable_weights": true,
     "tfr.ext.tfrbert.TFRBertRankingNetwork.transform": true,
+    "tfr.ext.tfrbert.TFRBertRankingNetwork.variable_dtype": true,
     "tfr.ext.tfrbert.TFRBertRankingNetwork.weights": true,
     "tfr.ext.tfrbert.TFRBertRankingNetwork.with_name_scope": true,
     "tfr.ext.tfrbert.TFRBertUtil": false,
@@ -989,6 +1283,7 @@
     "tfr.keras.canned.DNNRankingNetwork.add_weight": true,
     "tfr.keras.canned.DNNRankingNetwork.build": true,
     "tfr.keras.canned.DNNRankingNetwork.call": true,
+    "tfr.keras.canned.DNNRankingNetwork.compute_dtype": true,
     "tfr.keras.canned.DNNRankingNetwork.compute_logits": true,
     "tfr.keras.canned.DNNRankingNetwork.compute_mask": true,
     "tfr.keras.canned.DNNRankingNetwork.compute_output_shape": true,
@@ -996,6 +1291,7 @@
     "tfr.keras.canned.DNNRankingNetwork.context_feature_columns": true,
     "tfr.keras.canned.DNNRankingNetwork.count_params": true,
     "tfr.keras.canned.DNNRankingNetwork.dtype": true,
+    "tfr.keras.canned.DNNRankingNetwork.dtype_policy": true,
     "tfr.keras.canned.DNNRankingNetwork.dynamic": true,
     "tfr.keras.canned.DNNRankingNetwork.example_feature_columns": true,
     "tfr.keras.canned.DNNRankingNetwork.from_config": true,
@@ -1016,6 +1312,7 @@
     "tfr.keras.canned.DNNRankingNetwork.trainable": true,
     "tfr.keras.canned.DNNRankingNetwork.trainable_weights": true,
     "tfr.keras.canned.DNNRankingNetwork.transform": true,
+    "tfr.keras.canned.DNNRankingNetwork.variable_dtype": true,
     "tfr.keras.canned.DNNRankingNetwork.weights": true,
     "tfr.keras.canned.DNNRankingNetwork.with_name_scope": true,
     "tfr.keras.canned.GAMRankingNetwork": false,
@@ -1034,6 +1331,7 @@
     "tfr.keras.canned.GAMRankingNetwork.add_weight": true,
     "tfr.keras.canned.GAMRankingNetwork.build": true,
     "tfr.keras.canned.GAMRankingNetwork.call": true,
+    "tfr.keras.canned.GAMRankingNetwork.compute_dtype": true,
     "tfr.keras.canned.GAMRankingNetwork.compute_logits": true,
     "tfr.keras.canned.GAMRankingNetwork.compute_mask": true,
     "tfr.keras.canned.GAMRankingNetwork.compute_output_shape": true,
@@ -1041,6 +1339,7 @@
     "tfr.keras.canned.GAMRankingNetwork.context_feature_columns": true,
     "tfr.keras.canned.GAMRankingNetwork.count_params": true,
     "tfr.keras.canned.GAMRankingNetwork.dtype": true,
+    "tfr.keras.canned.GAMRankingNetwork.dtype_policy": true,
     "tfr.keras.canned.GAMRankingNetwork.dynamic": true,
     "tfr.keras.canned.GAMRankingNetwork.example_feature_columns": true,
     "tfr.keras.canned.GAMRankingNetwork.from_config": true,
@@ -1061,6 +1360,7 @@
     "tfr.keras.canned.GAMRankingNetwork.trainable": true,
     "tfr.keras.canned.GAMRankingNetwork.trainable_weights": true,
     "tfr.keras.canned.GAMRankingNetwork.transform": true,
+    "tfr.keras.canned.GAMRankingNetwork.variable_dtype": true,
     "tfr.keras.canned.GAMRankingNetwork.weights": true,
     "tfr.keras.canned.GAMRankingNetwork.with_name_scope": true,
     "tfr.keras.canned.dnn": false,
@@ -1080,6 +1380,7 @@
     "tfr.keras.canned.dnn.DNNRankingNetwork.add_weight": true,
     "tfr.keras.canned.dnn.DNNRankingNetwork.build": true,
     "tfr.keras.canned.dnn.DNNRankingNetwork.call": true,
+    "tfr.keras.canned.dnn.DNNRankingNetwork.compute_dtype": true,
     "tfr.keras.canned.dnn.DNNRankingNetwork.compute_logits": true,
     "tfr.keras.canned.dnn.DNNRankingNetwork.compute_mask": true,
     "tfr.keras.canned.dnn.DNNRankingNetwork.compute_output_shape": true,
@@ -1087,6 +1388,7 @@
     "tfr.keras.canned.dnn.DNNRankingNetwork.context_feature_columns": true,
     "tfr.keras.canned.dnn.DNNRankingNetwork.count_params": true,
     "tfr.keras.canned.dnn.DNNRankingNetwork.dtype": true,
+    "tfr.keras.canned.dnn.DNNRankingNetwork.dtype_policy": true,
     "tfr.keras.canned.dnn.DNNRankingNetwork.dynamic": true,
     "tfr.keras.canned.dnn.DNNRankingNetwork.example_feature_columns": true,
     "tfr.keras.canned.dnn.DNNRankingNetwork.from_config": true,
@@ -1107,6 +1409,7 @@
     "tfr.keras.canned.dnn.DNNRankingNetwork.trainable": true,
     "tfr.keras.canned.dnn.DNNRankingNetwork.trainable_weights": true,
     "tfr.keras.canned.dnn.DNNRankingNetwork.transform": true,
+    "tfr.keras.canned.dnn.DNNRankingNetwork.variable_dtype": true,
     "tfr.keras.canned.dnn.DNNRankingNetwork.weights": true,
     "tfr.keras.canned.dnn.DNNRankingNetwork.with_name_scope": true,
     "tfr.keras.canned.gam": false,
@@ -1126,6 +1429,7 @@
     "tfr.keras.canned.gam.GAMRankingNetwork.add_weight": true,
     "tfr.keras.canned.gam.GAMRankingNetwork.build": true,
     "tfr.keras.canned.gam.GAMRankingNetwork.call": true,
+    "tfr.keras.canned.gam.GAMRankingNetwork.compute_dtype": true,
     "tfr.keras.canned.gam.GAMRankingNetwork.compute_logits": true,
     "tfr.keras.canned.gam.GAMRankingNetwork.compute_mask": true,
     "tfr.keras.canned.gam.GAMRankingNetwork.compute_output_shape": true,
@@ -1133,6 +1437,7 @@
     "tfr.keras.canned.gam.GAMRankingNetwork.context_feature_columns": true,
     "tfr.keras.canned.gam.GAMRankingNetwork.count_params": true,
     "tfr.keras.canned.gam.GAMRankingNetwork.dtype": true,
+    "tfr.keras.canned.gam.GAMRankingNetwork.dtype_policy": true,
     "tfr.keras.canned.gam.GAMRankingNetwork.dynamic": true,
     "tfr.keras.canned.gam.GAMRankingNetwork.example_feature_columns": true,
     "tfr.keras.canned.gam.GAMRankingNetwork.from_config": true,
@@ -1153,6 +1458,7 @@
     "tfr.keras.canned.gam.GAMRankingNetwork.trainable": true,
     "tfr.keras.canned.gam.GAMRankingNetwork.trainable_weights": true,
     "tfr.keras.canned.gam.GAMRankingNetwork.transform": true,
+    "tfr.keras.canned.gam.GAMRankingNetwork.variable_dtype": true,
     "tfr.keras.canned.gam.GAMRankingNetwork.weights": true,
     "tfr.keras.canned.gam.GAMRankingNetwork.with_name_scope": true,
     "tfr.keras.estimator": false,
@@ -1174,12 +1480,14 @@
     "tfr.keras.feature.EncodeListwiseFeatures.add_weight": true,
     "tfr.keras.feature.EncodeListwiseFeatures.build": true,
     "tfr.keras.feature.EncodeListwiseFeatures.call": true,
+    "tfr.keras.feature.EncodeListwiseFeatures.compute_dtype": true,
     "tfr.keras.feature.EncodeListwiseFeatures.compute_mask": true,
     "tfr.keras.feature.EncodeListwiseFeatures.compute_output_shape": true,
     "tfr.keras.feature.EncodeListwiseFeatures.compute_output_signature": true,
     "tfr.keras.feature.EncodeListwiseFeatures.context_feature_columns": true,
     "tfr.keras.feature.EncodeListwiseFeatures.count_params": true,
     "tfr.keras.feature.EncodeListwiseFeatures.dtype": true,
+    "tfr.keras.feature.EncodeListwiseFeatures.dtype_policy": true,
     "tfr.keras.feature.EncodeListwiseFeatures.dynamic": true,
     "tfr.keras.feature.EncodeListwiseFeatures.example_feature_columns": true,
     "tfr.keras.feature.EncodeListwiseFeatures.from_config": true,
@@ -1198,6 +1506,7 @@
     "tfr.keras.feature.EncodeListwiseFeatures.supports_masking": true,
     "tfr.keras.feature.EncodeListwiseFeatures.trainable": true,
     "tfr.keras.feature.EncodeListwiseFeatures.trainable_weights": true,
+    "tfr.keras.feature.EncodeListwiseFeatures.variable_dtype": true,
     "tfr.keras.feature.EncodeListwiseFeatures.weights": true,
     "tfr.keras.feature.EncodeListwiseFeatures.with_name_scope": true,
     "tfr.keras.feature.GenerateMask": false,
@@ -1216,11 +1525,13 @@
     "tfr.keras.feature.GenerateMask.add_weight": true,
     "tfr.keras.feature.GenerateMask.build": true,
     "tfr.keras.feature.GenerateMask.call": true,
+    "tfr.keras.feature.GenerateMask.compute_dtype": true,
     "tfr.keras.feature.GenerateMask.compute_mask": true,
     "tfr.keras.feature.GenerateMask.compute_output_shape": true,
     "tfr.keras.feature.GenerateMask.compute_output_signature": true,
     "tfr.keras.feature.GenerateMask.count_params": true,
     "tfr.keras.feature.GenerateMask.dtype": true,
+    "tfr.keras.feature.GenerateMask.dtype_policy": true,
     "tfr.keras.feature.GenerateMask.dynamic": true,
     "tfr.keras.feature.GenerateMask.example_feature_columns": true,
     "tfr.keras.feature.GenerateMask.from_config": true,
@@ -1240,11 +1551,185 @@
     "tfr.keras.feature.GenerateMask.supports_masking": true,
     "tfr.keras.feature.GenerateMask.trainable": true,
     "tfr.keras.feature.GenerateMask.trainable_weights": true,
+    "tfr.keras.feature.GenerateMask.variable_dtype": true,
     "tfr.keras.feature.GenerateMask.weights": true,
     "tfr.keras.feature.GenerateMask.with_name_scope": true,
     "tfr.keras.feature.create_keras_inputs": false,
     "tfr.keras.feature.deserialize_feature_columns": false,
     "tfr.keras.feature.serialize_feature_columns": false,
+    "tfr.keras.layers": false,
+    "tfr.keras.layers.ConcatFeatures": false,
+    "tfr.keras.layers.ConcatFeatures.__call__": true,
+    "tfr.keras.layers.ConcatFeatures.__eq__": true,
+    "tfr.keras.layers.ConcatFeatures.__ge__": true,
+    "tfr.keras.layers.ConcatFeatures.__gt__": true,
+    "tfr.keras.layers.ConcatFeatures.__init__": true,
+    "tfr.keras.layers.ConcatFeatures.__le__": true,
+    "tfr.keras.layers.ConcatFeatures.__lt__": true,
+    "tfr.keras.layers.ConcatFeatures.__ne__": true,
+    "tfr.keras.layers.ConcatFeatures.__new__": true,
+    "tfr.keras.layers.ConcatFeatures.activity_regularizer": true,
+    "tfr.keras.layers.ConcatFeatures.add_loss": true,
+    "tfr.keras.layers.ConcatFeatures.add_metric": true,
+    "tfr.keras.layers.ConcatFeatures.add_weight": true,
+    "tfr.keras.layers.ConcatFeatures.build": true,
+    "tfr.keras.layers.ConcatFeatures.call": true,
+    "tfr.keras.layers.ConcatFeatures.compute_dtype": true,
+    "tfr.keras.layers.ConcatFeatures.compute_mask": true,
+    "tfr.keras.layers.ConcatFeatures.compute_output_shape": true,
+    "tfr.keras.layers.ConcatFeatures.compute_output_signature": true,
+    "tfr.keras.layers.ConcatFeatures.count_params": true,
+    "tfr.keras.layers.ConcatFeatures.dtype": true,
+    "tfr.keras.layers.ConcatFeatures.dtype_policy": true,
+    "tfr.keras.layers.ConcatFeatures.dynamic": true,
+    "tfr.keras.layers.ConcatFeatures.from_config": true,
+    "tfr.keras.layers.ConcatFeatures.get_config": true,
+    "tfr.keras.layers.ConcatFeatures.get_weights": true,
+    "tfr.keras.layers.ConcatFeatures.input": true,
+    "tfr.keras.layers.ConcatFeatures.input_spec": true,
+    "tfr.keras.layers.ConcatFeatures.losses": true,
+    "tfr.keras.layers.ConcatFeatures.metrics": true,
+    "tfr.keras.layers.ConcatFeatures.name": true,
+    "tfr.keras.layers.ConcatFeatures.name_scope": true,
+    "tfr.keras.layers.ConcatFeatures.non_trainable_weights": true,
+    "tfr.keras.layers.ConcatFeatures.output": true,
+    "tfr.keras.layers.ConcatFeatures.set_weights": true,
+    "tfr.keras.layers.ConcatFeatures.submodules": true,
+    "tfr.keras.layers.ConcatFeatures.supports_masking": true,
+    "tfr.keras.layers.ConcatFeatures.trainable": true,
+    "tfr.keras.layers.ConcatFeatures.trainable_weights": true,
+    "tfr.keras.layers.ConcatFeatures.variable_dtype": true,
+    "tfr.keras.layers.ConcatFeatures.weights": true,
+    "tfr.keras.layers.ConcatFeatures.with_name_scope": true,
+    "tfr.keras.layers.DocumentInteractionAttention": false,
+    "tfr.keras.layers.DocumentInteractionAttention.__call__": true,
+    "tfr.keras.layers.DocumentInteractionAttention.__eq__": true,
+    "tfr.keras.layers.DocumentInteractionAttention.__ge__": true,
+    "tfr.keras.layers.DocumentInteractionAttention.__gt__": true,
+    "tfr.keras.layers.DocumentInteractionAttention.__init__": true,
+    "tfr.keras.layers.DocumentInteractionAttention.__le__": true,
+    "tfr.keras.layers.DocumentInteractionAttention.__lt__": true,
+    "tfr.keras.layers.DocumentInteractionAttention.__ne__": true,
+    "tfr.keras.layers.DocumentInteractionAttention.__new__": true,
+    "tfr.keras.layers.DocumentInteractionAttention.activity_regularizer": true,
+    "tfr.keras.layers.DocumentInteractionAttention.add_loss": true,
+    "tfr.keras.layers.DocumentInteractionAttention.add_metric": true,
+    "tfr.keras.layers.DocumentInteractionAttention.add_weight": true,
+    "tfr.keras.layers.DocumentInteractionAttention.build": true,
+    "tfr.keras.layers.DocumentInteractionAttention.call": true,
+    "tfr.keras.layers.DocumentInteractionAttention.compute_dtype": true,
+    "tfr.keras.layers.DocumentInteractionAttention.compute_mask": true,
+    "tfr.keras.layers.DocumentInteractionAttention.compute_output_shape": true,
+    "tfr.keras.layers.DocumentInteractionAttention.compute_output_signature": true,
+    "tfr.keras.layers.DocumentInteractionAttention.count_params": true,
+    "tfr.keras.layers.DocumentInteractionAttention.dtype": true,
+    "tfr.keras.layers.DocumentInteractionAttention.dtype_policy": true,
+    "tfr.keras.layers.DocumentInteractionAttention.dynamic": true,
+    "tfr.keras.layers.DocumentInteractionAttention.from_config": true,
+    "tfr.keras.layers.DocumentInteractionAttention.get_config": true,
+    "tfr.keras.layers.DocumentInteractionAttention.get_weights": true,
+    "tfr.keras.layers.DocumentInteractionAttention.input": true,
+    "tfr.keras.layers.DocumentInteractionAttention.input_spec": true,
+    "tfr.keras.layers.DocumentInteractionAttention.losses": true,
+    "tfr.keras.layers.DocumentInteractionAttention.metrics": true,
+    "tfr.keras.layers.DocumentInteractionAttention.name": true,
+    "tfr.keras.layers.DocumentInteractionAttention.name_scope": true,
+    "tfr.keras.layers.DocumentInteractionAttention.non_trainable_weights": true,
+    "tfr.keras.layers.DocumentInteractionAttention.output": true,
+    "tfr.keras.layers.DocumentInteractionAttention.set_weights": true,
+    "tfr.keras.layers.DocumentInteractionAttention.submodules": true,
+    "tfr.keras.layers.DocumentInteractionAttention.supports_masking": true,
+    "tfr.keras.layers.DocumentInteractionAttention.trainable": true,
+    "tfr.keras.layers.DocumentInteractionAttention.trainable_weights": true,
+    "tfr.keras.layers.DocumentInteractionAttention.variable_dtype": true,
+    "tfr.keras.layers.DocumentInteractionAttention.weights": true,
+    "tfr.keras.layers.DocumentInteractionAttention.with_name_scope": true,
+    "tfr.keras.layers.FlattenList": false,
+    "tfr.keras.layers.FlattenList.__call__": true,
+    "tfr.keras.layers.FlattenList.__eq__": true,
+    "tfr.keras.layers.FlattenList.__ge__": true,
+    "tfr.keras.layers.FlattenList.__gt__": true,
+    "tfr.keras.layers.FlattenList.__init__": true,
+    "tfr.keras.layers.FlattenList.__le__": true,
+    "tfr.keras.layers.FlattenList.__lt__": true,
+    "tfr.keras.layers.FlattenList.__ne__": true,
+    "tfr.keras.layers.FlattenList.__new__": true,
+    "tfr.keras.layers.FlattenList.activity_regularizer": true,
+    "tfr.keras.layers.FlattenList.add_loss": true,
+    "tfr.keras.layers.FlattenList.add_metric": true,
+    "tfr.keras.layers.FlattenList.add_weight": true,
+    "tfr.keras.layers.FlattenList.build": true,
+    "tfr.keras.layers.FlattenList.call": true,
+    "tfr.keras.layers.FlattenList.compute_dtype": true,
+    "tfr.keras.layers.FlattenList.compute_mask": true,
+    "tfr.keras.layers.FlattenList.compute_output_shape": true,
+    "tfr.keras.layers.FlattenList.compute_output_signature": true,
+    "tfr.keras.layers.FlattenList.count_params": true,
+    "tfr.keras.layers.FlattenList.dtype": true,
+    "tfr.keras.layers.FlattenList.dtype_policy": true,
+    "tfr.keras.layers.FlattenList.dynamic": true,
+    "tfr.keras.layers.FlattenList.from_config": true,
+    "tfr.keras.layers.FlattenList.get_config": true,
+    "tfr.keras.layers.FlattenList.get_weights": true,
+    "tfr.keras.layers.FlattenList.input": true,
+    "tfr.keras.layers.FlattenList.input_spec": true,
+    "tfr.keras.layers.FlattenList.losses": true,
+    "tfr.keras.layers.FlattenList.metrics": true,
+    "tfr.keras.layers.FlattenList.name": true,
+    "tfr.keras.layers.FlattenList.name_scope": true,
+    "tfr.keras.layers.FlattenList.non_trainable_weights": true,
+    "tfr.keras.layers.FlattenList.output": true,
+    "tfr.keras.layers.FlattenList.set_weights": true,
+    "tfr.keras.layers.FlattenList.submodules": true,
+    "tfr.keras.layers.FlattenList.supports_masking": true,
+    "tfr.keras.layers.FlattenList.trainable": true,
+    "tfr.keras.layers.FlattenList.trainable_weights": true,
+    "tfr.keras.layers.FlattenList.variable_dtype": true,
+    "tfr.keras.layers.FlattenList.weights": true,
+    "tfr.keras.layers.FlattenList.with_name_scope": true,
+    "tfr.keras.layers.RestoreList": false,
+    "tfr.keras.layers.RestoreList.__call__": true,
+    "tfr.keras.layers.RestoreList.__eq__": true,
+    "tfr.keras.layers.RestoreList.__ge__": true,
+    "tfr.keras.layers.RestoreList.__gt__": true,
+    "tfr.keras.layers.RestoreList.__init__": true,
+    "tfr.keras.layers.RestoreList.__le__": true,
+    "tfr.keras.layers.RestoreList.__lt__": true,
+    "tfr.keras.layers.RestoreList.__ne__": true,
+    "tfr.keras.layers.RestoreList.__new__": true,
+    "tfr.keras.layers.RestoreList.activity_regularizer": true,
+    "tfr.keras.layers.RestoreList.add_loss": true,
+    "tfr.keras.layers.RestoreList.add_metric": true,
+    "tfr.keras.layers.RestoreList.add_weight": true,
+    "tfr.keras.layers.RestoreList.build": true,
+    "tfr.keras.layers.RestoreList.call": true,
+    "tfr.keras.layers.RestoreList.compute_dtype": true,
+    "tfr.keras.layers.RestoreList.compute_mask": true,
+    "tfr.keras.layers.RestoreList.compute_output_shape": true,
+    "tfr.keras.layers.RestoreList.compute_output_signature": true,
+    "tfr.keras.layers.RestoreList.count_params": true,
+    "tfr.keras.layers.RestoreList.dtype": true,
+    "tfr.keras.layers.RestoreList.dtype_policy": true,
+    "tfr.keras.layers.RestoreList.dynamic": true,
+    "tfr.keras.layers.RestoreList.from_config": true,
+    "tfr.keras.layers.RestoreList.get_config": true,
+    "tfr.keras.layers.RestoreList.get_weights": true,
+    "tfr.keras.layers.RestoreList.input": true,
+    "tfr.keras.layers.RestoreList.input_spec": true,
+    "tfr.keras.layers.RestoreList.losses": true,
+    "tfr.keras.layers.RestoreList.metrics": true,
+    "tfr.keras.layers.RestoreList.name": true,
+    "tfr.keras.layers.RestoreList.name_scope": true,
+    "tfr.keras.layers.RestoreList.non_trainable_weights": true,
+    "tfr.keras.layers.RestoreList.output": true,
+    "tfr.keras.layers.RestoreList.set_weights": true,
+    "tfr.keras.layers.RestoreList.submodules": true,
+    "tfr.keras.layers.RestoreList.supports_masking": true,
+    "tfr.keras.layers.RestoreList.trainable": true,
+    "tfr.keras.layers.RestoreList.trainable_weights": true,
+    "tfr.keras.layers.RestoreList.variable_dtype": true,
+    "tfr.keras.layers.RestoreList.weights": true,
+    "tfr.keras.layers.RestoreList.with_name_scope": true,
     "tfr.keras.losses": false,
     "tfr.keras.losses.ApproxMRRLoss": false,
     "tfr.keras.losses.ApproxMRRLoss.__call__": true,
@@ -1272,6 +1757,31 @@
     "tfr.keras.losses.ApproxNDCGLoss.call": true,
     "tfr.keras.losses.ApproxNDCGLoss.from_config": true,
     "tfr.keras.losses.ApproxNDCGLoss.get_config": true,
+    "tfr.keras.losses.ClickEMLoss": false,
+    "tfr.keras.losses.ClickEMLoss.__call__": true,
+    "tfr.keras.losses.ClickEMLoss.__eq__": true,
+    "tfr.keras.losses.ClickEMLoss.__ge__": true,
+    "tfr.keras.losses.ClickEMLoss.__gt__": true,
+    "tfr.keras.losses.ClickEMLoss.__init__": true,
+    "tfr.keras.losses.ClickEMLoss.__le__": true,
+    "tfr.keras.losses.ClickEMLoss.__lt__": true,
+    "tfr.keras.losses.ClickEMLoss.__ne__": true,
+    "tfr.keras.losses.ClickEMLoss.__new__": true,
+    "tfr.keras.losses.ClickEMLoss.call": true,
+    "tfr.keras.losses.ClickEMLoss.from_config": true,
+    "tfr.keras.losses.ClickEMLoss.get_config": true,
+    "tfr.keras.losses.DCGLambdaWeight": false,
+    "tfr.keras.losses.DCGLambdaWeight.__eq__": true,
+    "tfr.keras.losses.DCGLambdaWeight.__ge__": true,
+    "tfr.keras.losses.DCGLambdaWeight.__gt__": true,
+    "tfr.keras.losses.DCGLambdaWeight.__init__": true,
+    "tfr.keras.losses.DCGLambdaWeight.__le__": true,
+    "tfr.keras.losses.DCGLambdaWeight.__lt__": true,
+    "tfr.keras.losses.DCGLambdaWeight.__ne__": true,
+    "tfr.keras.losses.DCGLambdaWeight.__new__": true,
+    "tfr.keras.losses.DCGLambdaWeight.get_config": true,
+    "tfr.keras.losses.DCGLambdaWeight.individual_weights": true,
+    "tfr.keras.losses.DCGLambdaWeight.pair_weights": true,
     "tfr.keras.losses.GumbelApproxNDCGLoss": false,
     "tfr.keras.losses.GumbelApproxNDCGLoss.__call__": true,
     "tfr.keras.losses.GumbelApproxNDCGLoss.__eq__": true,
@@ -1285,6 +1795,18 @@
     "tfr.keras.losses.GumbelApproxNDCGLoss.call": true,
     "tfr.keras.losses.GumbelApproxNDCGLoss.from_config": true,
     "tfr.keras.losses.GumbelApproxNDCGLoss.get_config": true,
+    "tfr.keras.losses.ListMLELambdaWeight": false,
+    "tfr.keras.losses.ListMLELambdaWeight.__eq__": true,
+    "tfr.keras.losses.ListMLELambdaWeight.__ge__": true,
+    "tfr.keras.losses.ListMLELambdaWeight.__gt__": true,
+    "tfr.keras.losses.ListMLELambdaWeight.__init__": true,
+    "tfr.keras.losses.ListMLELambdaWeight.__le__": true,
+    "tfr.keras.losses.ListMLELambdaWeight.__lt__": true,
+    "tfr.keras.losses.ListMLELambdaWeight.__ne__": true,
+    "tfr.keras.losses.ListMLELambdaWeight.__new__": true,
+    "tfr.keras.losses.ListMLELambdaWeight.get_config": true,
+    "tfr.keras.losses.ListMLELambdaWeight.individual_weights": true,
+    "tfr.keras.losses.ListMLELambdaWeight.pair_weights": true,
     "tfr.keras.losses.ListMLELoss": false,
     "tfr.keras.losses.ListMLELoss.__call__": true,
     "tfr.keras.losses.ListMLELoss.__eq__": true,
@@ -1311,6 +1833,18 @@
     "tfr.keras.losses.MeanSquaredLoss.call": true,
     "tfr.keras.losses.MeanSquaredLoss.from_config": true,
     "tfr.keras.losses.MeanSquaredLoss.get_config": true,
+    "tfr.keras.losses.NDCGLambdaWeight": false,
+    "tfr.keras.losses.NDCGLambdaWeight.__eq__": true,
+    "tfr.keras.losses.NDCGLambdaWeight.__ge__": true,
+    "tfr.keras.losses.NDCGLambdaWeight.__gt__": true,
+    "tfr.keras.losses.NDCGLambdaWeight.__init__": true,
+    "tfr.keras.losses.NDCGLambdaWeight.__le__": true,
+    "tfr.keras.losses.NDCGLambdaWeight.__lt__": true,
+    "tfr.keras.losses.NDCGLambdaWeight.__ne__": true,
+    "tfr.keras.losses.NDCGLambdaWeight.__new__": true,
+    "tfr.keras.losses.NDCGLambdaWeight.get_config": true,
+    "tfr.keras.losses.NDCGLambdaWeight.individual_weights": true,
+    "tfr.keras.losses.NDCGLambdaWeight.pair_weights": true,
     "tfr.keras.losses.PairwiseHingeLoss": false,
     "tfr.keras.losses.PairwiseHingeLoss.__call__": true,
     "tfr.keras.losses.PairwiseHingeLoss.__eq__": true,
@@ -1350,6 +1884,18 @@
     "tfr.keras.losses.PairwiseSoftZeroOneLoss.call": true,
     "tfr.keras.losses.PairwiseSoftZeroOneLoss.from_config": true,
     "tfr.keras.losses.PairwiseSoftZeroOneLoss.get_config": true,
+    "tfr.keras.losses.PrecisionLambdaWeight": false,
+    "tfr.keras.losses.PrecisionLambdaWeight.__eq__": true,
+    "tfr.keras.losses.PrecisionLambdaWeight.__ge__": true,
+    "tfr.keras.losses.PrecisionLambdaWeight.__gt__": true,
+    "tfr.keras.losses.PrecisionLambdaWeight.__init__": true,
+    "tfr.keras.losses.PrecisionLambdaWeight.__le__": true,
+    "tfr.keras.losses.PrecisionLambdaWeight.__lt__": true,
+    "tfr.keras.losses.PrecisionLambdaWeight.__ne__": true,
+    "tfr.keras.losses.PrecisionLambdaWeight.__new__": true,
+    "tfr.keras.losses.PrecisionLambdaWeight.get_config": true,
+    "tfr.keras.losses.PrecisionLambdaWeight.individual_weights": true,
+    "tfr.keras.losses.PrecisionLambdaWeight.pair_weights": true,
     "tfr.keras.losses.RankingLossKey": false,
     "tfr.keras.losses.RankingLossKey.APPROX_MRR_LOSS": true,
     "tfr.keras.losses.RankingLossKey.APPROX_NDCG_LOSS": true,
@@ -1409,10 +1955,7 @@
     "tfr.keras.losses.UniqueSoftmaxLoss.call": true,
     "tfr.keras.losses.UniqueSoftmaxLoss.from_config": true,
     "tfr.keras.losses.UniqueSoftmaxLoss.get_config": true,
-    "tfr.keras.losses.absolute_import": true,
-    "tfr.keras.losses.division": true,
     "tfr.keras.losses.get": false,
-    "tfr.keras.losses.print_function": true,
     "tfr.keras.metrics": false,
     "tfr.keras.metrics.ARPMetric": false,
     "tfr.keras.metrics.ARPMetric.__call__": true,
@@ -1430,11 +1973,13 @@
     "tfr.keras.metrics.ARPMetric.add_weight": true,
     "tfr.keras.metrics.ARPMetric.build": true,
     "tfr.keras.metrics.ARPMetric.call": true,
+    "tfr.keras.metrics.ARPMetric.compute_dtype": true,
     "tfr.keras.metrics.ARPMetric.compute_mask": true,
     "tfr.keras.metrics.ARPMetric.compute_output_shape": true,
     "tfr.keras.metrics.ARPMetric.compute_output_signature": true,
     "tfr.keras.metrics.ARPMetric.count_params": true,
     "tfr.keras.metrics.ARPMetric.dtype": true,
+    "tfr.keras.metrics.ARPMetric.dtype_policy": true,
     "tfr.keras.metrics.ARPMetric.dynamic": true,
     "tfr.keras.metrics.ARPMetric.from_config": true,
     "tfr.keras.metrics.ARPMetric.get_config": true,
@@ -1455,6 +2000,7 @@
     "tfr.keras.metrics.ARPMetric.trainable": true,
     "tfr.keras.metrics.ARPMetric.trainable_weights": true,
     "tfr.keras.metrics.ARPMetric.update_state": true,
+    "tfr.keras.metrics.ARPMetric.variable_dtype": true,
     "tfr.keras.metrics.ARPMetric.weights": true,
     "tfr.keras.metrics.ARPMetric.with_name_scope": true,
     "tfr.keras.metrics.AlphaDCGMetric": false,
@@ -1473,11 +2019,13 @@
     "tfr.keras.metrics.AlphaDCGMetric.add_weight": true,
     "tfr.keras.metrics.AlphaDCGMetric.build": true,
     "tfr.keras.metrics.AlphaDCGMetric.call": true,
+    "tfr.keras.metrics.AlphaDCGMetric.compute_dtype": true,
     "tfr.keras.metrics.AlphaDCGMetric.compute_mask": true,
     "tfr.keras.metrics.AlphaDCGMetric.compute_output_shape": true,
     "tfr.keras.metrics.AlphaDCGMetric.compute_output_signature": true,
     "tfr.keras.metrics.AlphaDCGMetric.count_params": true,
     "tfr.keras.metrics.AlphaDCGMetric.dtype": true,
+    "tfr.keras.metrics.AlphaDCGMetric.dtype_policy": true,
     "tfr.keras.metrics.AlphaDCGMetric.dynamic": true,
     "tfr.keras.metrics.AlphaDCGMetric.from_config": true,
     "tfr.keras.metrics.AlphaDCGMetric.get_config": true,
@@ -1498,6 +2046,7 @@
     "tfr.keras.metrics.AlphaDCGMetric.trainable": true,
     "tfr.keras.metrics.AlphaDCGMetric.trainable_weights": true,
     "tfr.keras.metrics.AlphaDCGMetric.update_state": true,
+    "tfr.keras.metrics.AlphaDCGMetric.variable_dtype": true,
     "tfr.keras.metrics.AlphaDCGMetric.weights": true,
     "tfr.keras.metrics.AlphaDCGMetric.with_name_scope": true,
     "tfr.keras.metrics.DCGMetric": false,
@@ -1516,11 +2065,13 @@
     "tfr.keras.metrics.DCGMetric.add_weight": true,
     "tfr.keras.metrics.DCGMetric.build": true,
     "tfr.keras.metrics.DCGMetric.call": true,
+    "tfr.keras.metrics.DCGMetric.compute_dtype": true,
     "tfr.keras.metrics.DCGMetric.compute_mask": true,
     "tfr.keras.metrics.DCGMetric.compute_output_shape": true,
     "tfr.keras.metrics.DCGMetric.compute_output_signature": true,
     "tfr.keras.metrics.DCGMetric.count_params": true,
     "tfr.keras.metrics.DCGMetric.dtype": true,
+    "tfr.keras.metrics.DCGMetric.dtype_policy": true,
     "tfr.keras.metrics.DCGMetric.dynamic": true,
     "tfr.keras.metrics.DCGMetric.from_config": true,
     "tfr.keras.metrics.DCGMetric.get_config": true,
@@ -1541,6 +2092,7 @@
     "tfr.keras.metrics.DCGMetric.trainable": true,
     "tfr.keras.metrics.DCGMetric.trainable_weights": true,
     "tfr.keras.metrics.DCGMetric.update_state": true,
+    "tfr.keras.metrics.DCGMetric.variable_dtype": true,
     "tfr.keras.metrics.DCGMetric.weights": true,
     "tfr.keras.metrics.DCGMetric.with_name_scope": true,
     "tfr.keras.metrics.MRRMetric": false,
@@ -1559,11 +2111,13 @@
     "tfr.keras.metrics.MRRMetric.add_weight": true,
     "tfr.keras.metrics.MRRMetric.build": true,
     "tfr.keras.metrics.MRRMetric.call": true,
+    "tfr.keras.metrics.MRRMetric.compute_dtype": true,
     "tfr.keras.metrics.MRRMetric.compute_mask": true,
     "tfr.keras.metrics.MRRMetric.compute_output_shape": true,
     "tfr.keras.metrics.MRRMetric.compute_output_signature": true,
     "tfr.keras.metrics.MRRMetric.count_params": true,
     "tfr.keras.metrics.MRRMetric.dtype": true,
+    "tfr.keras.metrics.MRRMetric.dtype_policy": true,
     "tfr.keras.metrics.MRRMetric.dynamic": true,
     "tfr.keras.metrics.MRRMetric.from_config": true,
     "tfr.keras.metrics.MRRMetric.get_config": true,
@@ -1584,6 +2138,7 @@
     "tfr.keras.metrics.MRRMetric.trainable": true,
     "tfr.keras.metrics.MRRMetric.trainable_weights": true,
     "tfr.keras.metrics.MRRMetric.update_state": true,
+    "tfr.keras.metrics.MRRMetric.variable_dtype": true,
     "tfr.keras.metrics.MRRMetric.weights": true,
     "tfr.keras.metrics.MRRMetric.with_name_scope": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric": false,
@@ -1602,11 +2157,13 @@
     "tfr.keras.metrics.MeanAveragePrecisionMetric.add_weight": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.build": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.call": true,
+    "tfr.keras.metrics.MeanAveragePrecisionMetric.compute_dtype": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.compute_mask": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.compute_output_shape": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.compute_output_signature": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.count_params": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.dtype": true,
+    "tfr.keras.metrics.MeanAveragePrecisionMetric.dtype_policy": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.dynamic": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.from_config": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.get_config": true,
@@ -1627,6 +2184,7 @@
     "tfr.keras.metrics.MeanAveragePrecisionMetric.trainable": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.trainable_weights": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.update_state": true,
+    "tfr.keras.metrics.MeanAveragePrecisionMetric.variable_dtype": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.weights": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.with_name_scope": true,
     "tfr.keras.metrics.NDCGMetric": false,
@@ -1645,11 +2203,13 @@
     "tfr.keras.metrics.NDCGMetric.add_weight": true,
     "tfr.keras.metrics.NDCGMetric.build": true,
     "tfr.keras.metrics.NDCGMetric.call": true,
+    "tfr.keras.metrics.NDCGMetric.compute_dtype": true,
     "tfr.keras.metrics.NDCGMetric.compute_mask": true,
     "tfr.keras.metrics.NDCGMetric.compute_output_shape": true,
     "tfr.keras.metrics.NDCGMetric.compute_output_signature": true,
     "tfr.keras.metrics.NDCGMetric.count_params": true,
     "tfr.keras.metrics.NDCGMetric.dtype": true,
+    "tfr.keras.metrics.NDCGMetric.dtype_policy": true,
     "tfr.keras.metrics.NDCGMetric.dynamic": true,
     "tfr.keras.metrics.NDCGMetric.from_config": true,
     "tfr.keras.metrics.NDCGMetric.get_config": true,
@@ -1670,6 +2230,7 @@
     "tfr.keras.metrics.NDCGMetric.trainable": true,
     "tfr.keras.metrics.NDCGMetric.trainable_weights": true,
     "tfr.keras.metrics.NDCGMetric.update_state": true,
+    "tfr.keras.metrics.NDCGMetric.variable_dtype": true,
     "tfr.keras.metrics.NDCGMetric.weights": true,
     "tfr.keras.metrics.NDCGMetric.with_name_scope": true,
     "tfr.keras.metrics.OPAMetric": false,
@@ -1688,11 +2249,13 @@
     "tfr.keras.metrics.OPAMetric.add_weight": true,
     "tfr.keras.metrics.OPAMetric.build": true,
     "tfr.keras.metrics.OPAMetric.call": true,
+    "tfr.keras.metrics.OPAMetric.compute_dtype": true,
     "tfr.keras.metrics.OPAMetric.compute_mask": true,
     "tfr.keras.metrics.OPAMetric.compute_output_shape": true,
     "tfr.keras.metrics.OPAMetric.compute_output_signature": true,
     "tfr.keras.metrics.OPAMetric.count_params": true,
     "tfr.keras.metrics.OPAMetric.dtype": true,
+    "tfr.keras.metrics.OPAMetric.dtype_policy": true,
     "tfr.keras.metrics.OPAMetric.dynamic": true,
     "tfr.keras.metrics.OPAMetric.from_config": true,
     "tfr.keras.metrics.OPAMetric.get_config": true,
@@ -1713,6 +2276,7 @@
     "tfr.keras.metrics.OPAMetric.trainable": true,
     "tfr.keras.metrics.OPAMetric.trainable_weights": true,
     "tfr.keras.metrics.OPAMetric.update_state": true,
+    "tfr.keras.metrics.OPAMetric.variable_dtype": true,
     "tfr.keras.metrics.OPAMetric.weights": true,
     "tfr.keras.metrics.OPAMetric.with_name_scope": true,
     "tfr.keras.metrics.PrecisionIAMetric": false,
@@ -1731,11 +2295,13 @@
     "tfr.keras.metrics.PrecisionIAMetric.add_weight": true,
     "tfr.keras.metrics.PrecisionIAMetric.build": true,
     "tfr.keras.metrics.PrecisionIAMetric.call": true,
+    "tfr.keras.metrics.PrecisionIAMetric.compute_dtype": true,
     "tfr.keras.metrics.PrecisionIAMetric.compute_mask": true,
     "tfr.keras.metrics.PrecisionIAMetric.compute_output_shape": true,
     "tfr.keras.metrics.PrecisionIAMetric.compute_output_signature": true,
     "tfr.keras.metrics.PrecisionIAMetric.count_params": true,
     "tfr.keras.metrics.PrecisionIAMetric.dtype": true,
+    "tfr.keras.metrics.PrecisionIAMetric.dtype_policy": true,
     "tfr.keras.metrics.PrecisionIAMetric.dynamic": true,
     "tfr.keras.metrics.PrecisionIAMetric.from_config": true,
     "tfr.keras.metrics.PrecisionIAMetric.get_config": true,
@@ -1756,6 +2322,7 @@
     "tfr.keras.metrics.PrecisionIAMetric.trainable": true,
     "tfr.keras.metrics.PrecisionIAMetric.trainable_weights": true,
     "tfr.keras.metrics.PrecisionIAMetric.update_state": true,
+    "tfr.keras.metrics.PrecisionIAMetric.variable_dtype": true,
     "tfr.keras.metrics.PrecisionIAMetric.weights": true,
     "tfr.keras.metrics.PrecisionIAMetric.with_name_scope": true,
     "tfr.keras.metrics.PrecisionMetric": false,
@@ -1774,11 +2341,13 @@
     "tfr.keras.metrics.PrecisionMetric.add_weight": true,
     "tfr.keras.metrics.PrecisionMetric.build": true,
     "tfr.keras.metrics.PrecisionMetric.call": true,
+    "tfr.keras.metrics.PrecisionMetric.compute_dtype": true,
     "tfr.keras.metrics.PrecisionMetric.compute_mask": true,
     "tfr.keras.metrics.PrecisionMetric.compute_output_shape": true,
     "tfr.keras.metrics.PrecisionMetric.compute_output_signature": true,
     "tfr.keras.metrics.PrecisionMetric.count_params": true,
     "tfr.keras.metrics.PrecisionMetric.dtype": true,
+    "tfr.keras.metrics.PrecisionMetric.dtype_policy": true,
     "tfr.keras.metrics.PrecisionMetric.dynamic": true,
     "tfr.keras.metrics.PrecisionMetric.from_config": true,
     "tfr.keras.metrics.PrecisionMetric.get_config": true,
@@ -1799,11 +2368,77 @@
     "tfr.keras.metrics.PrecisionMetric.trainable": true,
     "tfr.keras.metrics.PrecisionMetric.trainable_weights": true,
     "tfr.keras.metrics.PrecisionMetric.update_state": true,
+    "tfr.keras.metrics.PrecisionMetric.variable_dtype": true,
     "tfr.keras.metrics.PrecisionMetric.weights": true,
     "tfr.keras.metrics.PrecisionMetric.with_name_scope": true,
+    "tfr.keras.metrics.RankingMetricKey": false,
+    "tfr.keras.metrics.RankingMetricKey.ALPHA_DCG": true,
+    "tfr.keras.metrics.RankingMetricKey.ARP": true,
+    "tfr.keras.metrics.RankingMetricKey.DCG": true,
+    "tfr.keras.metrics.RankingMetricKey.MAP": true,
+    "tfr.keras.metrics.RankingMetricKey.MRR": true,
+    "tfr.keras.metrics.RankingMetricKey.NDCG": true,
+    "tfr.keras.metrics.RankingMetricKey.ORDERED_PAIR_ACCURACY": true,
+    "tfr.keras.metrics.RankingMetricKey.PRECISION": true,
+    "tfr.keras.metrics.RankingMetricKey.PRECISION_IA": true,
+    "tfr.keras.metrics.RankingMetricKey.__eq__": true,
+    "tfr.keras.metrics.RankingMetricKey.__ge__": true,
+    "tfr.keras.metrics.RankingMetricKey.__gt__": true,
+    "tfr.keras.metrics.RankingMetricKey.__init__": true,
+    "tfr.keras.metrics.RankingMetricKey.__le__": true,
+    "tfr.keras.metrics.RankingMetricKey.__lt__": true,
+    "tfr.keras.metrics.RankingMetricKey.__ne__": true,
+    "tfr.keras.metrics.RankingMetricKey.__new__": true,
+    "tfr.keras.metrics.RecallMetric": false,
+    "tfr.keras.metrics.RecallMetric.__call__": true,
+    "tfr.keras.metrics.RecallMetric.__eq__": true,
+    "tfr.keras.metrics.RecallMetric.__ge__": true,
+    "tfr.keras.metrics.RecallMetric.__gt__": true,
+    "tfr.keras.metrics.RecallMetric.__init__": true,
+    "tfr.keras.metrics.RecallMetric.__le__": true,
+    "tfr.keras.metrics.RecallMetric.__lt__": true,
+    "tfr.keras.metrics.RecallMetric.__ne__": true,
+    "tfr.keras.metrics.RecallMetric.__new__": true,
+    "tfr.keras.metrics.RecallMetric.activity_regularizer": true,
+    "tfr.keras.metrics.RecallMetric.add_loss": true,
+    "tfr.keras.metrics.RecallMetric.add_metric": true,
+    "tfr.keras.metrics.RecallMetric.add_weight": true,
+    "tfr.keras.metrics.RecallMetric.build": true,
+    "tfr.keras.metrics.RecallMetric.call": true,
+    "tfr.keras.metrics.RecallMetric.compute_dtype": true,
+    "tfr.keras.metrics.RecallMetric.compute_mask": true,
+    "tfr.keras.metrics.RecallMetric.compute_output_shape": true,
+    "tfr.keras.metrics.RecallMetric.compute_output_signature": true,
+    "tfr.keras.metrics.RecallMetric.count_params": true,
+    "tfr.keras.metrics.RecallMetric.dtype": true,
+    "tfr.keras.metrics.RecallMetric.dtype_policy": true,
+    "tfr.keras.metrics.RecallMetric.dynamic": true,
+    "tfr.keras.metrics.RecallMetric.from_config": true,
+    "tfr.keras.metrics.RecallMetric.get_config": true,
+    "tfr.keras.metrics.RecallMetric.get_weights": true,
+    "tfr.keras.metrics.RecallMetric.input": true,
+    "tfr.keras.metrics.RecallMetric.input_spec": true,
+    "tfr.keras.metrics.RecallMetric.losses": true,
+    "tfr.keras.metrics.RecallMetric.metrics": true,
+    "tfr.keras.metrics.RecallMetric.name": true,
+    "tfr.keras.metrics.RecallMetric.name_scope": true,
+    "tfr.keras.metrics.RecallMetric.non_trainable_weights": true,
+    "tfr.keras.metrics.RecallMetric.output": true,
+    "tfr.keras.metrics.RecallMetric.reset_states": true,
+    "tfr.keras.metrics.RecallMetric.result": true,
+    "tfr.keras.metrics.RecallMetric.set_weights": true,
+    "tfr.keras.metrics.RecallMetric.submodules": true,
+    "tfr.keras.metrics.RecallMetric.supports_masking": true,
+    "tfr.keras.metrics.RecallMetric.trainable": true,
+    "tfr.keras.metrics.RecallMetric.trainable_weights": true,
+    "tfr.keras.metrics.RecallMetric.update_state": true,
+    "tfr.keras.metrics.RecallMetric.variable_dtype": true,
+    "tfr.keras.metrics.RecallMetric.weights": true,
+    "tfr.keras.metrics.RecallMetric.with_name_scope": true,
     "tfr.keras.metrics.absolute_import": true,
     "tfr.keras.metrics.default_keras_metrics": false,
     "tfr.keras.metrics.division": true,
+    "tfr.keras.metrics.get": false,
     "tfr.keras.metrics.print_function": true,
     "tfr.keras.model": false,
     "tfr.keras.model.create_keras_model": false,
@@ -1824,6 +2459,7 @@
     "tfr.keras.network.RankingNetwork.add_weight": true,
     "tfr.keras.network.RankingNetwork.build": true,
     "tfr.keras.network.RankingNetwork.call": true,
+    "tfr.keras.network.RankingNetwork.compute_dtype": true,
     "tfr.keras.network.RankingNetwork.compute_logits": true,
     "tfr.keras.network.RankingNetwork.compute_mask": true,
     "tfr.keras.network.RankingNetwork.compute_output_shape": true,
@@ -1831,6 +2467,7 @@
     "tfr.keras.network.RankingNetwork.context_feature_columns": true,
     "tfr.keras.network.RankingNetwork.count_params": true,
     "tfr.keras.network.RankingNetwork.dtype": true,
+    "tfr.keras.network.RankingNetwork.dtype_policy": true,
     "tfr.keras.network.RankingNetwork.dynamic": true,
     "tfr.keras.network.RankingNetwork.example_feature_columns": true,
     "tfr.keras.network.RankingNetwork.from_config": true,
@@ -1850,6 +2487,7 @@
     "tfr.keras.network.RankingNetwork.trainable": true,
     "tfr.keras.network.RankingNetwork.trainable_weights": true,
     "tfr.keras.network.RankingNetwork.transform": true,
+    "tfr.keras.network.RankingNetwork.variable_dtype": true,
     "tfr.keras.network.RankingNetwork.weights": true,
     "tfr.keras.network.RankingNetwork.with_name_scope": true,
     "tfr.keras.network.UnivariateRankingNetwork": false,
@@ -1868,6 +2506,7 @@
     "tfr.keras.network.UnivariateRankingNetwork.add_weight": true,
     "tfr.keras.network.UnivariateRankingNetwork.build": true,
     "tfr.keras.network.UnivariateRankingNetwork.call": true,
+    "tfr.keras.network.UnivariateRankingNetwork.compute_dtype": true,
     "tfr.keras.network.UnivariateRankingNetwork.compute_logits": true,
     "tfr.keras.network.UnivariateRankingNetwork.compute_mask": true,
     "tfr.keras.network.UnivariateRankingNetwork.compute_output_shape": true,
@@ -1875,6 +2514,7 @@
     "tfr.keras.network.UnivariateRankingNetwork.context_feature_columns": true,
     "tfr.keras.network.UnivariateRankingNetwork.count_params": true,
     "tfr.keras.network.UnivariateRankingNetwork.dtype": true,
+    "tfr.keras.network.UnivariateRankingNetwork.dtype_policy": true,
     "tfr.keras.network.UnivariateRankingNetwork.dynamic": true,
     "tfr.keras.network.UnivariateRankingNetwork.example_feature_columns": true,
     "tfr.keras.network.UnivariateRankingNetwork.from_config": true,
@@ -1895,9 +2535,36 @@
     "tfr.keras.network.UnivariateRankingNetwork.trainable": true,
     "tfr.keras.network.UnivariateRankingNetwork.trainable_weights": true,
     "tfr.keras.network.UnivariateRankingNetwork.transform": true,
+    "tfr.keras.network.UnivariateRankingNetwork.variable_dtype": true,
     "tfr.keras.network.UnivariateRankingNetwork.weights": true,
     "tfr.keras.network.UnivariateRankingNetwork.with_name_scope": true,
     "tfr.keras.network.listwise_scoring": false,
+    "tfr.keras.saved_model": false,
+    "tfr.keras.saved_model.Signatures": false,
+    "tfr.keras.saved_model.Signatures.__call__": true,
+    "tfr.keras.saved_model.Signatures.__eq__": true,
+    "tfr.keras.saved_model.Signatures.__ge__": true,
+    "tfr.keras.saved_model.Signatures.__gt__": true,
+    "tfr.keras.saved_model.Signatures.__init__": true,
+    "tfr.keras.saved_model.Signatures.__le__": true,
+    "tfr.keras.saved_model.Signatures.__lt__": true,
+    "tfr.keras.saved_model.Signatures.__ne__": true,
+    "tfr.keras.saved_model.Signatures.__new__": true,
+    "tfr.keras.saved_model.Signatures.name": true,
+    "tfr.keras.saved_model.Signatures.name_scope": true,
+    "tfr.keras.saved_model.Signatures.non_trainable_variables": true,
+    "tfr.keras.saved_model.Signatures.predict_tf_function": true,
+    "tfr.keras.saved_model.Signatures.regress_tf_function": true,
+    "tfr.keras.saved_model.Signatures.submodules": true,
+    "tfr.keras.saved_model.Signatures.trainable_variables": true,
+    "tfr.keras.saved_model.Signatures.variables": true,
+    "tfr.keras.saved_model.Signatures.with_name_scope": true,
+    "tfr.keras.utils": false,
+    "tfr.keras.utils.identity": false,
+    "tfr.keras.utils.inverse": false,
+    "tfr.keras.utils.is_greater_equal_1": false,
+    "tfr.keras.utils.log2_inverse": false,
+    "tfr.keras.utils.pow_minus_1": false,
     "tfr.losses": false,
     "tfr.losses.RankingLossKey": false,
     "tfr.losses.RankingLossKey.APPROX_MRR_LOSS": true,
@@ -1933,6 +2600,7 @@
     "tfr.metrics.RankingMetricKey": false,
     "tfr.metrics.RankingMetricKey.ALPHA_DCG": true,
     "tfr.metrics.RankingMetricKey.ARP": true,
+    "tfr.metrics.RankingMetricKey.BPREF": true,
     "tfr.metrics.RankingMetricKey.DCG": true,
     "tfr.metrics.RankingMetricKey.MAP": true,
     "tfr.metrics.RankingMetricKey.MRR": true,
@@ -1951,6 +2619,7 @@
     "tfr.metrics.absolute_import": true,
     "tfr.metrics.alpha_discounted_cumulative_gain": false,
     "tfr.metrics.average_relevance_position": false,
+    "tfr.metrics.binary_preference": false,
     "tfr.metrics.compute_mean": false,
     "tfr.metrics.discounted_cumulative_gain": false,
     "tfr.metrics.division": true,
@@ -1972,10 +2641,12 @@
     "tfr.utils.absolute_import": true,
     "tfr.utils.de_noise": false,
     "tfr.utils.division": true,
+    "tfr.utils.gather_per_row": false,
     "tfr.utils.is_label_valid": false,
     "tfr.utils.organize_valid_indices": false,
     "tfr.utils.padded_nd_indices": false,
     "tfr.utils.print_function": true,
+    "tfr.utils.ragged_to_dense": false,
     "tfr.utils.reshape_first_ndims": false,
     "tfr.utils.reshape_to_2d": false,
     "tfr.utils.shuffle_valid_indices": false,
@@ -1984,5 +2655,6 @@
   },
   "py_module_names": [
     "tfr"
-  ]
+  ],
+  "site_link": null
 }
diff --git a/tensorflow_ranking/g3doc/api_docs/python/_redirects.yaml b/tensorflow_ranking/g3doc/api_docs/python/tfr/_redirects.yaml
similarity index 100%
rename from tensorflow_ranking/g3doc/api_docs/python/_redirects.yaml
rename to tensorflow_ranking/g3doc/api_docs/python/tfr/_redirects.yaml
diff --git a/tensorflow_ranking/g3doc/api_docs/python/_toc.yaml b/tensorflow_ranking/g3doc/api_docs/python/tfr/_toc.yaml
similarity index 81%
rename from tensorflow_ranking/g3doc/api_docs/python/_toc.yaml
rename to tensorflow_ranking/g3doc/api_docs/python/tfr/_toc.yaml
index 091ab09..13f6e0b 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/_toc.yaml
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/_toc.yaml
@@ -17,6 +17,8 @@ toc:
     path: /ranking/api_docs/python/tfr/data/build_ranking_serving_input_receiver_fn_with_parsing_fn
   - title: build_sequence_example_serving_input_receiver_fn
     path: /ranking/api_docs/python/tfr/data/build_sequence_example_serving_input_receiver_fn
+  - title: build_tf_example_serving_input_receiver_fn
+    path: /ranking/api_docs/python/tfr/data/build_tf_example_serving_input_receiver_fn
   - title: make_parsing_fn
     path: /ranking/api_docs/python/tfr/data/make_parsing_fn
   - title: parse_from_example_in_example
@@ -25,6 +27,8 @@ toc:
     path: /ranking/api_docs/python/tfr/data/parse_from_example_list
   - title: parse_from_sequence_example
     path: /ranking/api_docs/python/tfr/data/parse_from_sequence_example
+  - title: parse_from_tf_example
+    path: /ranking/api_docs/python/tfr/data/parse_from_tf_example
   - title: read_batched_sequence_example_dataset
     path: /ranking/api_docs/python/tfr/data/read_batched_sequence_example_dataset
 - title: tfr.estimator
@@ -107,16 +111,28 @@ toc:
     section:
     - title: Overview
       path: /ranking/api_docs/python/tfr/keras/feature
-    - title: create_keras_inputs
-      path: /ranking/api_docs/python/tfr/keras/feature/create_keras_inputs
-    - title: deserialize_feature_columns
-      path: /ranking/api_docs/python/tfr/keras/feature/deserialize_feature_columns
     - title: EncodeListwiseFeatures
       path: /ranking/api_docs/python/tfr/keras/feature/EncodeListwiseFeatures
     - title: GenerateMask
       path: /ranking/api_docs/python/tfr/keras/feature/GenerateMask
+    - title: create_keras_inputs
+      path: /ranking/api_docs/python/tfr/keras/feature/create_keras_inputs
+    - title: deserialize_feature_columns
+      path: /ranking/api_docs/python/tfr/keras/feature/deserialize_feature_columns
     - title: serialize_feature_columns
       path: /ranking/api_docs/python/tfr/keras/feature/serialize_feature_columns
+  - title: layers
+    section:
+    - title: Overview
+      path: /ranking/api_docs/python/tfr/keras/layers
+    - title: ConcatFeatures
+      path: /ranking/api_docs/python/tfr/keras/layers/ConcatFeatures
+    - title: DocumentInteractionAttention
+      path: /ranking/api_docs/python/tfr/keras/layers/DocumentInteractionAttention
+    - title: FlattenList
+      path: /ranking/api_docs/python/tfr/keras/layers/FlattenList
+    - title: RestoreList
+      path: /ranking/api_docs/python/tfr/keras/layers/RestoreList
   - title: losses
     section:
     - title: Overview
@@ -125,20 +141,28 @@ toc:
       path: /ranking/api_docs/python/tfr/keras/losses/ApproxMRRLoss
     - title: ApproxNDCGLoss
       path: /ranking/api_docs/python/tfr/keras/losses/ApproxNDCGLoss
-    - title: get
-      path: /ranking/api_docs/python/tfr/keras/losses/get
+    - title: ClickEMLoss
+      path: /ranking/api_docs/python/tfr/keras/losses/ClickEMLoss
+    - title: DCGLambdaWeight
+      path: /ranking/api_docs/python/tfr/keras/losses/DCGLambdaWeight
     - title: GumbelApproxNDCGLoss
       path: /ranking/api_docs/python/tfr/keras/losses/GumbelApproxNDCGLoss
+    - title: ListMLELambdaWeight
+      path: /ranking/api_docs/python/tfr/keras/losses/ListMLELambdaWeight
     - title: ListMLELoss
       path: /ranking/api_docs/python/tfr/keras/losses/ListMLELoss
     - title: MeanSquaredLoss
       path: /ranking/api_docs/python/tfr/keras/losses/MeanSquaredLoss
+    - title: NDCGLambdaWeight
+      path: /ranking/api_docs/python/tfr/keras/losses/NDCGLambdaWeight
     - title: PairwiseHingeLoss
       path: /ranking/api_docs/python/tfr/keras/losses/PairwiseHingeLoss
     - title: PairwiseLogisticLoss
       path: /ranking/api_docs/python/tfr/keras/losses/PairwiseLogisticLoss
     - title: PairwiseSoftZeroOneLoss
       path: /ranking/api_docs/python/tfr/keras/losses/PairwiseSoftZeroOneLoss
+    - title: PrecisionLambdaWeight
+      path: /ranking/api_docs/python/tfr/keras/losses/PrecisionLambdaWeight
     - title: RankingLossKey
       path: /ranking/api_docs/python/tfr/keras/losses/RankingLossKey
     - title: SigmoidCrossEntropyLoss
@@ -147,22 +171,22 @@ toc:
       path: /ranking/api_docs/python/tfr/keras/losses/SoftmaxLoss
     - title: UniqueSoftmaxLoss
       path: /ranking/api_docs/python/tfr/keras/losses/UniqueSoftmaxLoss
+    - title: get
+      path: /ranking/api_docs/python/tfr/keras/losses/get
   - title: metrics
     section:
     - title: Overview
       path: /ranking/api_docs/python/tfr/keras/metrics
-    - title: AlphaDCGMetric
-      path: /ranking/api_docs/python/tfr/keras/metrics/AlphaDCGMetric
     - title: ARPMetric
       path: /ranking/api_docs/python/tfr/keras/metrics/ARPMetric
+    - title: AlphaDCGMetric
+      path: /ranking/api_docs/python/tfr/keras/metrics/AlphaDCGMetric
     - title: DCGMetric
       path: /ranking/api_docs/python/tfr/keras/metrics/DCGMetric
-    - title: default_keras_metrics
-      path: /ranking/api_docs/python/tfr/keras/metrics/default_keras_metrics
-    - title: MeanAveragePrecisionMetric
-      path: /ranking/api_docs/python/tfr/keras/metrics/MeanAveragePrecisionMetric
     - title: MRRMetric
       path: /ranking/api_docs/python/tfr/keras/metrics/MRRMetric
+    - title: MeanAveragePrecisionMetric
+      path: /ranking/api_docs/python/tfr/keras/metrics/MeanAveragePrecisionMetric
     - title: NDCGMetric
       path: /ranking/api_docs/python/tfr/keras/metrics/NDCGMetric
     - title: OPAMetric
@@ -171,6 +195,14 @@ toc:
       path: /ranking/api_docs/python/tfr/keras/metrics/PrecisionIAMetric
     - title: PrecisionMetric
       path: /ranking/api_docs/python/tfr/keras/metrics/PrecisionMetric
+    - title: RankingMetricKey
+      path: /ranking/api_docs/python/tfr/keras/metrics/RankingMetricKey
+    - title: RecallMetric
+      path: /ranking/api_docs/python/tfr/keras/metrics/RecallMetric
+    - title: default_keras_metrics
+      path: /ranking/api_docs/python/tfr/keras/metrics/default_keras_metrics
+    - title: get
+      path: /ranking/api_docs/python/tfr/keras/metrics/get
   - title: model
     section:
     - title: Overview
@@ -181,16 +213,38 @@ toc:
     section:
     - title: Overview
       path: /ranking/api_docs/python/tfr/keras/network
-    - title: listwise_scoring
-      path: /ranking/api_docs/python/tfr/keras/network/listwise_scoring
     - title: RankingNetwork
       path: /ranking/api_docs/python/tfr/keras/network/RankingNetwork
     - title: UnivariateRankingNetwork
       path: /ranking/api_docs/python/tfr/keras/network/UnivariateRankingNetwork
+    - title: listwise_scoring
+      path: /ranking/api_docs/python/tfr/keras/network/listwise_scoring
+  - title: saved_model
+    section:
+    - title: Overview
+      path: /ranking/api_docs/python/tfr/keras/saved_model
+    - title: Signatures
+      path: /ranking/api_docs/python/tfr/keras/saved_model/Signatures
+  - title: utils
+    section:
+    - title: Overview
+      path: /ranking/api_docs/python/tfr/keras/utils
+    - title: identity
+      path: /ranking/api_docs/python/tfr/keras/utils/identity
+    - title: inverse
+      path: /ranking/api_docs/python/tfr/keras/utils/inverse
+    - title: is_greater_equal_1
+      path: /ranking/api_docs/python/tfr/keras/utils/is_greater_equal_1
+    - title: log2_inverse
+      path: /ranking/api_docs/python/tfr/keras/utils/log2_inverse
+    - title: pow_minus_1
+      path: /ranking/api_docs/python/tfr/keras/utils/pow_minus_1
 - title: tfr.losses
   section:
   - title: Overview
     path: /ranking/api_docs/python/tfr/losses
+  - title: RankingLossKey
+    path: /ranking/api_docs/python/tfr/losses/RankingLossKey
   - title: create_ndcg_lambda_weight
     path: /ranking/api_docs/python/tfr/losses/create_ndcg_lambda_weight
   - title: create_p_list_mle_lambda_weight
@@ -201,16 +255,18 @@ toc:
     path: /ranking/api_docs/python/tfr/losses/make_loss_fn
   - title: make_loss_metric_fn
     path: /ranking/api_docs/python/tfr/losses/make_loss_metric_fn
-  - title: RankingLossKey
-    path: /ranking/api_docs/python/tfr/losses/RankingLossKey
 - title: tfr.metrics
   section:
   - title: Overview
     path: /ranking/api_docs/python/tfr/metrics
+  - title: RankingMetricKey
+    path: /ranking/api_docs/python/tfr/metrics/RankingMetricKey
   - title: alpha_discounted_cumulative_gain
     path: /ranking/api_docs/python/tfr/metrics/alpha_discounted_cumulative_gain
   - title: average_relevance_position
     path: /ranking/api_docs/python/tfr/metrics/average_relevance_position
+  - title: binary_preference
+    path: /ranking/api_docs/python/tfr/metrics/binary_preference
   - title: compute_mean
     path: /ranking/api_docs/python/tfr/metrics/compute_mean
   - title: discounted_cumulative_gain
@@ -231,8 +287,6 @@ toc:
     path: /ranking/api_docs/python/tfr/metrics/precision
   - title: precision_ia
     path: /ranking/api_docs/python/tfr/metrics/precision_ia
-  - title: RankingMetricKey
-    path: /ranking/api_docs/python/tfr/metrics/RankingMetricKey
 - title: tfr.model
   section:
   - title: Overview
@@ -245,19 +299,23 @@ toc:
     path: /ranking/api_docs/python/tfr/utils
   - title: de_noise
     path: /ranking/api_docs/python/tfr/utils/de_noise
+  - title: gather_per_row
+    path: /ranking/api_docs/python/tfr/utils/gather_per_row
   - title: is_label_valid
     path: /ranking/api_docs/python/tfr/utils/is_label_valid
   - title: organize_valid_indices
     path: /ranking/api_docs/python/tfr/utils/organize_valid_indices
   - title: padded_nd_indices
     path: /ranking/api_docs/python/tfr/utils/padded_nd_indices
+  - title: ragged_to_dense
+    path: /ranking/api_docs/python/tfr/utils/ragged_to_dense
   - title: reshape_first_ndims
     path: /ranking/api_docs/python/tfr/utils/reshape_first_ndims
   - title: reshape_to_2d
     path: /ranking/api_docs/python/tfr/utils/reshape_to_2d
   - title: shuffle_valid_indices
     path: /ranking/api_docs/python/tfr/utils/shuffle_valid_indices
-  - title: sorted_ranks
-    path: /ranking/api_docs/python/tfr/utils/sorted_ranks
   - title: sort_by_scores
     path: /ranking/api_docs/python/tfr/utils/sort_by_scores
+  - title: sorted_ranks
+    path: /ranking/api_docs/python/tfr/utils/sorted_ranks
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/all_symbols.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/all_symbols.md
new file mode 100644
index 0000000..c4d4ba0
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/all_symbols.md
@@ -0,0 +1,143 @@
+# All symbols in TensorFlow Ranking
+
+<!-- Insert buttons and diff -->
+
+## Primary symbols
+
+*   <a href="../tfr.md"><code>tfr</code></a>
+*   <a href="../tfr/data.md"><code>tfr.data</code></a>
+*   <a href="../tfr/data/build_ranking_dataset.md"><code>tfr.data.build_ranking_dataset</code></a>
+*   <a href="../tfr/data/build_ranking_dataset_with_parsing_fn.md"><code>tfr.data.build_ranking_dataset_with_parsing_fn</code></a>
+*   <a href="../tfr/data/build_ranking_serving_input_receiver_fn.md"><code>tfr.data.build_ranking_serving_input_receiver_fn</code></a>
+*   <a href="../tfr/data/build_ranking_serving_input_receiver_fn_with_parsing_fn.md"><code>tfr.data.build_ranking_serving_input_receiver_fn_with_parsing_fn</code></a>
+*   <a href="../tfr/data/build_sequence_example_serving_input_receiver_fn.md"><code>tfr.data.build_sequence_example_serving_input_receiver_fn</code></a>
+*   <a href="../tfr/data/build_tf_example_serving_input_receiver_fn.md"><code>tfr.data.build_tf_example_serving_input_receiver_fn</code></a>
+*   <a href="../tfr/data/make_parsing_fn.md"><code>tfr.data.make_parsing_fn</code></a>
+*   <a href="../tfr/data/parse_from_example_in_example.md"><code>tfr.data.parse_from_example_in_example</code></a>
+*   <a href="../tfr/data/parse_from_example_list.md"><code>tfr.data.parse_from_example_list</code></a>
+*   <a href="../tfr/data/parse_from_sequence_example.md"><code>tfr.data.parse_from_sequence_example</code></a>
+*   <a href="../tfr/data/parse_from_tf_example.md"><code>tfr.data.parse_from_tf_example</code></a>
+*   <a href="../tfr/data/read_batched_sequence_example_dataset.md"><code>tfr.data.read_batched_sequence_example_dataset</code></a>
+*   <a href="../tfr/estimator.md"><code>tfr.estimator</code></a>
+*   <a href="../tfr/estimator/EstimatorBuilder.md"><code>tfr.estimator.EstimatorBuilder</code></a>
+*   <a href="../tfr/estimator/GAMEstimatorBuilder.md"><code>tfr.estimator.GAMEstimatorBuilder</code></a>
+*   <a href="../tfr/estimator/make_dnn_ranking_estimator.md"><code>tfr.estimator.make_dnn_ranking_estimator</code></a>
+*   <a href="../tfr/estimator/make_gam_ranking_estimator.md"><code>tfr.estimator.make_gam_ranking_estimator</code></a>
+*   <a href="../tfr/ext.md"><code>tfr.ext</code></a>
+*   <a href="../tfr/ext/pipeline.md"><code>tfr.ext.pipeline</code></a>
+*   <a href="../tfr/ext/pipeline/RankingPipeline.md"><code>tfr.ext.pipeline.RankingPipeline</code></a>
+*   <a href="../tfr/ext/tfrbert.md"><code>tfr.ext.tfrbert</code></a>
+*   <a href="../tfr/ext/tfrbert/TFRBertRankingNetwork.md"><code>tfr.ext.tfrbert.TFRBertRankingNetwork</code></a>
+*   <a href="../tfr/ext/tfrbert/TFRBertUtil.md"><code>tfr.ext.tfrbert.TFRBertUtil</code></a>
+*   <a href="../tfr/feature.md"><code>tfr.feature</code></a>
+*   <a href="../tfr/feature/encode_features.md"><code>tfr.feature.encode_features</code></a>
+*   <a href="../tfr/feature/encode_listwise_features.md"><code>tfr.feature.encode_listwise_features</code></a>
+*   <a href="../tfr/feature/encode_pointwise_features.md"><code>tfr.feature.encode_pointwise_features</code></a>
+*   <a href="../tfr/feature/make_identity_transform_fn.md"><code>tfr.feature.make_identity_transform_fn</code></a>
+*   <a href="../tfr/head.md"><code>tfr.head</code></a>
+*   <a href="../tfr/head/create_multi_ranking_head.md"><code>tfr.head.create_multi_ranking_head</code></a>
+*   <a href="../tfr/head/create_ranking_head.md"><code>tfr.head.create_ranking_head</code></a>
+*   <a href="../tfr/keras.md"><code>tfr.keras</code></a>
+*   <a href="../tfr/keras/canned.md"><code>tfr.keras.canned</code></a>
+*   <a href="../tfr/keras/canned/DNNRankingNetwork.md"><code>tfr.keras.canned.DNNRankingNetwork</code></a>
+*   <a href="../tfr/keras/canned/GAMRankingNetwork.md"><code>tfr.keras.canned.GAMRankingNetwork</code></a>
+*   <a href="../tfr/keras/canned/dnn.md"><code>tfr.keras.canned.dnn</code></a>
+*   <a href="../tfr/keras/canned/DNNRankingNetwork.md"><code>tfr.keras.canned.dnn.DNNRankingNetwork</code></a>
+*   <a href="../tfr/keras/canned/gam.md"><code>tfr.keras.canned.gam</code></a>
+*   <a href="../tfr/keras/canned/GAMRankingNetwork.md"><code>tfr.keras.canned.gam.GAMRankingNetwork</code></a>
+*   <a href="../tfr/keras/estimator.md"><code>tfr.keras.estimator</code></a>
+*   <a href="../tfr/keras/estimator/model_to_estimator.md"><code>tfr.keras.estimator.model_to_estimator</code></a>
+*   <a href="../tfr/keras/feature.md"><code>tfr.keras.feature</code></a>
+*   <a href="../tfr/keras/feature/EncodeListwiseFeatures.md"><code>tfr.keras.feature.EncodeListwiseFeatures</code></a>
+*   <a href="../tfr/keras/feature/GenerateMask.md"><code>tfr.keras.feature.GenerateMask</code></a>
+*   <a href="../tfr/keras/feature/create_keras_inputs.md"><code>tfr.keras.feature.create_keras_inputs</code></a>
+*   <a href="../tfr/keras/feature/deserialize_feature_columns.md"><code>tfr.keras.feature.deserialize_feature_columns</code></a>
+*   <a href="../tfr/keras/feature/serialize_feature_columns.md"><code>tfr.keras.feature.serialize_feature_columns</code></a>
+*   <a href="../tfr/keras/layers.md"><code>tfr.keras.layers</code></a>
+*   <a href="../tfr/keras/layers/ConcatFeatures.md"><code>tfr.keras.layers.ConcatFeatures</code></a>
+*   <a href="../tfr/keras/layers/DocumentInteractionAttention.md"><code>tfr.keras.layers.DocumentInteractionAttention</code></a>
+*   <a href="../tfr/keras/layers/FlattenList.md"><code>tfr.keras.layers.FlattenList</code></a>
+*   <a href="../tfr/keras/layers/RestoreList.md"><code>tfr.keras.layers.RestoreList</code></a>
+*   <a href="../tfr/keras/losses.md"><code>tfr.keras.losses</code></a>
+*   <a href="../tfr/keras/losses/ApproxMRRLoss.md"><code>tfr.keras.losses.ApproxMRRLoss</code></a>
+*   <a href="../tfr/keras/losses/ApproxNDCGLoss.md"><code>tfr.keras.losses.ApproxNDCGLoss</code></a>
+*   <a href="../tfr/keras/losses/ClickEMLoss.md"><code>tfr.keras.losses.ClickEMLoss</code></a>
+*   <a href="../tfr/keras/losses/DCGLambdaWeight.md"><code>tfr.keras.losses.DCGLambdaWeight</code></a>
+*   <a href="../tfr/keras/losses/GumbelApproxNDCGLoss.md"><code>tfr.keras.losses.GumbelApproxNDCGLoss</code></a>
+*   <a href="../tfr/keras/losses/ListMLELambdaWeight.md"><code>tfr.keras.losses.ListMLELambdaWeight</code></a>
+*   <a href="../tfr/keras/losses/ListMLELoss.md"><code>tfr.keras.losses.ListMLELoss</code></a>
+*   <a href="../tfr/keras/losses/MeanSquaredLoss.md"><code>tfr.keras.losses.MeanSquaredLoss</code></a>
+*   <a href="../tfr/keras/losses/NDCGLambdaWeight.md"><code>tfr.keras.losses.NDCGLambdaWeight</code></a>
+*   <a href="../tfr/keras/losses/PairwiseHingeLoss.md"><code>tfr.keras.losses.PairwiseHingeLoss</code></a>
+*   <a href="../tfr/keras/losses/PairwiseLogisticLoss.md"><code>tfr.keras.losses.PairwiseLogisticLoss</code></a>
+*   <a href="../tfr/keras/losses/PairwiseSoftZeroOneLoss.md"><code>tfr.keras.losses.PairwiseSoftZeroOneLoss</code></a>
+*   <a href="../tfr/keras/losses/PrecisionLambdaWeight.md"><code>tfr.keras.losses.PrecisionLambdaWeight</code></a>
+*   <a href="../tfr/keras/losses/RankingLossKey.md"><code>tfr.keras.losses.RankingLossKey</code></a>
+*   <a href="../tfr/keras/losses/SigmoidCrossEntropyLoss.md"><code>tfr.keras.losses.SigmoidCrossEntropyLoss</code></a>
+*   <a href="../tfr/keras/losses/SoftmaxLoss.md"><code>tfr.keras.losses.SoftmaxLoss</code></a>
+*   <a href="../tfr/keras/losses/UniqueSoftmaxLoss.md"><code>tfr.keras.losses.UniqueSoftmaxLoss</code></a>
+*   <a href="../tfr/keras/losses/get.md"><code>tfr.keras.losses.get</code></a>
+*   <a href="../tfr/keras/metrics.md"><code>tfr.keras.metrics</code></a>
+*   <a href="../tfr/keras/metrics/ARPMetric.md"><code>tfr.keras.metrics.ARPMetric</code></a>
+*   <a href="../tfr/keras/metrics/AlphaDCGMetric.md"><code>tfr.keras.metrics.AlphaDCGMetric</code></a>
+*   <a href="../tfr/keras/metrics/DCGMetric.md"><code>tfr.keras.metrics.DCGMetric</code></a>
+*   <a href="../tfr/keras/metrics/MRRMetric.md"><code>tfr.keras.metrics.MRRMetric</code></a>
+*   <a href="../tfr/keras/metrics/MeanAveragePrecisionMetric.md"><code>tfr.keras.metrics.MeanAveragePrecisionMetric</code></a>
+*   <a href="../tfr/keras/metrics/NDCGMetric.md"><code>tfr.keras.metrics.NDCGMetric</code></a>
+*   <a href="../tfr/keras/metrics/OPAMetric.md"><code>tfr.keras.metrics.OPAMetric</code></a>
+*   <a href="../tfr/keras/metrics/PrecisionIAMetric.md"><code>tfr.keras.metrics.PrecisionIAMetric</code></a>
+*   <a href="../tfr/keras/metrics/PrecisionMetric.md"><code>tfr.keras.metrics.PrecisionMetric</code></a>
+*   <a href="../tfr/keras/metrics/RankingMetricKey.md"><code>tfr.keras.metrics.RankingMetricKey</code></a>
+*   <a href="../tfr/keras/metrics/RecallMetric.md"><code>tfr.keras.metrics.RecallMetric</code></a>
+*   <a href="../tfr/keras/metrics/default_keras_metrics.md"><code>tfr.keras.metrics.default_keras_metrics</code></a>
+*   <a href="../tfr/keras/metrics/get.md"><code>tfr.keras.metrics.get</code></a>
+*   <a href="../tfr/keras/model.md"><code>tfr.keras.model</code></a>
+*   <a href="../tfr/keras/model/create_keras_model.md"><code>tfr.keras.model.create_keras_model</code></a>
+*   <a href="../tfr/keras/network.md"><code>tfr.keras.network</code></a>
+*   <a href="../tfr/keras/network/RankingNetwork.md"><code>tfr.keras.network.RankingNetwork</code></a>
+*   <a href="../tfr/keras/network/UnivariateRankingNetwork.md"><code>tfr.keras.network.UnivariateRankingNetwork</code></a>
+*   <a href="../tfr/keras/network/listwise_scoring.md"><code>tfr.keras.network.listwise_scoring</code></a>
+*   <a href="../tfr/keras/saved_model.md"><code>tfr.keras.saved_model</code></a>
+*   <a href="../tfr/keras/saved_model/Signatures.md"><code>tfr.keras.saved_model.Signatures</code></a>
+*   <a href="../tfr/keras/utils.md"><code>tfr.keras.utils</code></a>
+*   <a href="../tfr/keras/utils/identity.md"><code>tfr.keras.utils.identity</code></a>
+*   <a href="../tfr/keras/utils/inverse.md"><code>tfr.keras.utils.inverse</code></a>
+*   <a href="../tfr/keras/utils/is_greater_equal_1.md"><code>tfr.keras.utils.is_greater_equal_1</code></a>
+*   <a href="../tfr/keras/utils/log2_inverse.md"><code>tfr.keras.utils.log2_inverse</code></a>
+*   <a href="../tfr/keras/utils/pow_minus_1.md"><code>tfr.keras.utils.pow_minus_1</code></a>
+*   <a href="../tfr/losses.md"><code>tfr.losses</code></a>
+*   <a href="../tfr/losses/RankingLossKey.md"><code>tfr.losses.RankingLossKey</code></a>
+*   <a href="../tfr/losses/create_ndcg_lambda_weight.md"><code>tfr.losses.create_ndcg_lambda_weight</code></a>
+*   <a href="../tfr/losses/create_p_list_mle_lambda_weight.md"><code>tfr.losses.create_p_list_mle_lambda_weight</code></a>
+*   <a href="../tfr/losses/create_reciprocal_rank_lambda_weight.md"><code>tfr.losses.create_reciprocal_rank_lambda_weight</code></a>
+*   <a href="../tfr/losses/make_loss_fn.md"><code>tfr.losses.make_loss_fn</code></a>
+*   <a href="../tfr/losses/make_loss_metric_fn.md"><code>tfr.losses.make_loss_metric_fn</code></a>
+*   <a href="../tfr/metrics.md"><code>tfr.metrics</code></a>
+*   <a href="../tfr/metrics/RankingMetricKey.md"><code>tfr.metrics.RankingMetricKey</code></a>
+*   <a href="../tfr/metrics/alpha_discounted_cumulative_gain.md"><code>tfr.metrics.alpha_discounted_cumulative_gain</code></a>
+*   <a href="../tfr/metrics/average_relevance_position.md"><code>tfr.metrics.average_relevance_position</code></a>
+*   <a href="../tfr/metrics/binary_preference.md"><code>tfr.metrics.binary_preference</code></a>
+*   <a href="../tfr/metrics/compute_mean.md"><code>tfr.metrics.compute_mean</code></a>
+*   <a href="../tfr/metrics/discounted_cumulative_gain.md"><code>tfr.metrics.discounted_cumulative_gain</code></a>
+*   <a href="../tfr/metrics/eval_metric.md"><code>tfr.metrics.eval_metric</code></a>
+*   <a href="../tfr/metrics/make_ranking_metric_fn.md"><code>tfr.metrics.make_ranking_metric_fn</code></a>
+*   <a href="../tfr/metrics/mean_average_precision.md"><code>tfr.metrics.mean_average_precision</code></a>
+*   <a href="../tfr/metrics/mean_reciprocal_rank.md"><code>tfr.metrics.mean_reciprocal_rank</code></a>
+*   <a href="../tfr/metrics/normalized_discounted_cumulative_gain.md"><code>tfr.metrics.normalized_discounted_cumulative_gain</code></a>
+*   <a href="../tfr/metrics/ordered_pair_accuracy.md"><code>tfr.metrics.ordered_pair_accuracy</code></a>
+*   <a href="../tfr/metrics/precision.md"><code>tfr.metrics.precision</code></a>
+*   <a href="../tfr/metrics/precision_ia.md"><code>tfr.metrics.precision_ia</code></a>
+*   <a href="../tfr/model.md"><code>tfr.model</code></a>
+*   <a href="../tfr/model/make_groupwise_ranking_fn.md"><code>tfr.model.make_groupwise_ranking_fn</code></a>
+*   <a href="../tfr/utils.md"><code>tfr.utils</code></a>
+*   <a href="../tfr/utils/de_noise.md"><code>tfr.utils.de_noise</code></a>
+*   <a href="../tfr/utils/gather_per_row.md"><code>tfr.utils.gather_per_row</code></a>
+*   <a href="../tfr/utils/is_label_valid.md"><code>tfr.utils.is_label_valid</code></a>
+*   <a href="../tfr/utils/organize_valid_indices.md"><code>tfr.utils.organize_valid_indices</code></a>
+*   <a href="../tfr/utils/padded_nd_indices.md"><code>tfr.utils.padded_nd_indices</code></a>
+*   <a href="../tfr/utils/ragged_to_dense.md"><code>tfr.utils.ragged_to_dense</code></a>
+*   <a href="../tfr/utils/reshape_first_ndims.md"><code>tfr.utils.reshape_first_ndims</code></a>
+*   <a href="../tfr/utils/reshape_to_2d.md"><code>tfr.utils.reshape_to_2d</code></a>
+*   <a href="../tfr/utils/shuffle_valid_indices.md"><code>tfr.utils.shuffle_valid_indices</code></a>
+*   <a href="../tfr/utils/sort_by_scores.md"><code>tfr.utils.sort_by_scores</code></a>
+*   <a href="../tfr/utils/sorted_ranks.md"><code>tfr.utils.sorted_ranks</code></a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/data.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/data.md
index efd585a..22ed0e1 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/data.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/data.md
@@ -1,3 +1,5 @@
+description: Input data parsing for tf-ranking library.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.data" />
 <meta itemprop="path" content="Stable" />
@@ -10,8 +12,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
@@ -42,6 +43,9 @@ Returns a receiver function with the provided `parsing_fn`.
 [`build_sequence_example_serving_input_receiver_fn(...)`](../tfr/data/build_sequence_example_serving_input_receiver_fn.md):
 Creates a serving_input_receiver_fn for `SequenceExample` inputs.
 
+[`build_tf_example_serving_input_receiver_fn(...)`](../tfr/data/build_tf_example_serving_input_receiver_fn.md):
+Builds a serving input fn for `tensorflow.training.Example`.
+
 [`make_parsing_fn(...)`](../tfr/data/make_parsing_fn.md): Returns a parsing fn
 for a standard data format.
 
@@ -54,11 +58,38 @@ an `ExampleListWithContext` batch to a feature map.
 [`parse_from_sequence_example(...)`](../tfr/data/parse_from_sequence_example.md):
 Parses SequenceExample to feature maps.
 
+[`parse_from_tf_example(...)`](../tfr/data/parse_from_tf_example.md): Parse
+function to convert `tf.train.Example` to feature maps.
+
 [`read_batched_sequence_example_dataset(...)`](../tfr/data/read_batched_sequence_example_dataset.md):
 Returns a `Dataset` of features from `SequenceExample`.
 
-## Other Members
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Other Members</h2></th></tr>
 
-*   `EIE = 'example_in_example'` <a id="EIE"></a>
-*   `ELWC = 'example_list_with_context'` <a id="ELWC"></a>
-*   `SEQ = 'sequence_example'` <a id="SEQ"></a>
+<tr>
+<td>
+EIE<a id="EIE"></a>
+</td>
+<td>
+`'example_in_example'`
+</td>
+</tr><tr>
+<td>
+ELWC<a id="ELWC"></a>
+</td>
+<td>
+`'example_list_with_context'`
+</td>
+</tr><tr>
+<td>
+SEQ<a id="SEQ"></a>
+</td>
+<td>
+`'sequence_example'`
+</td>
+</tr>
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_dataset.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_dataset.md
index 0d3a8ea..6da1b48 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_dataset.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_dataset.md
@@ -1,3 +1,5 @@
+description: Builds a ranking tf.dataset with a standard data format.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.data.build_ranking_dataset" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py#L954-L1002">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -23,7 +24,7 @@ Builds a ranking tf.dataset with a standard data format.
 <code>tfr.data.build_ranking_dataset(
     file_pattern, data_format, batch_size, context_feature_spec,
     example_feature_spec, list_size=None, size_feature_name=None,
-    shuffle_examples=False, seed=None, **kwargs
+    mask_feature_name=None, shuffle_examples=False, seed=None, **kwargs
 )
 </code></pre>
 
@@ -88,6 +89,16 @@ generated.
 </td>
 </tr><tr>
 <td>
+`mask_feature_name`
+</td>
+<td>
+(str) Name of feature for example list masks. Populates
+the feature dictionary with a `tf.bool` Tensor of shape [batch_size,
+list_size] for this feature name. If None, which is default, this feature
+is not generated.
+</td>
+</tr><tr>
+<td>
 `shuffle_examples`
 </td>
 <td>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_dataset_with_parsing_fn.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_dataset_with_parsing_fn.md
index ded9063..5cf1485 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_dataset_with_parsing_fn.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_dataset_with_parsing_fn.md
@@ -1,3 +1,5 @@
+description: Builds a ranking tf.dataset using the provided parsing_fn.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.data.build_ranking_dataset_with_parsing_fn" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py#L857-L951">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_serving_input_receiver_fn.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_serving_input_receiver_fn.md
index bfd0e1e..de0bf9a 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_serving_input_receiver_fn.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_serving_input_receiver_fn.md
@@ -1,3 +1,5 @@
+description: Returns a serving input receiver fn for a standard data format.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.data.build_ranking_serving_input_receiver_fn" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py#L1035-L1079">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -22,8 +23,8 @@ Returns a serving input receiver fn for a standard data format.
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.data.build_ranking_serving_input_receiver_fn(
     data_format, context_feature_spec, example_feature_spec, list_size=None,
-    size_feature_name=None, receiver_name='input_ranking_data',
-    default_batch_size=None
+    size_feature_name=None, mask_feature_name=None,
+    receiver_name=&#x27;input_ranking_data&#x27;, default_batch_size=None
 )
 </code></pre>
 
@@ -78,6 +79,16 @@ generated.
 </td>
 </tr><tr>
 <td>
+`mask_feature_name`
+</td>
+<td>
+(str) Name of feature for example list masks. Populates
+the feature dictionary with a `tf.bool` Tensor of shape [batch_size,
+list_size] for this feature name. If None, which is default, this feature
+is not generated.
+</td>
+</tr><tr>
+<td>
 `receiver_name`
 </td>
 <td>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_serving_input_receiver_fn_with_parsing_fn.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_serving_input_receiver_fn_with_parsing_fn.md
index 6129e79..7222804 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_serving_input_receiver_fn_with_parsing_fn.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_ranking_serving_input_receiver_fn_with_parsing_fn.md
@@ -1,3 +1,5 @@
+description: Returns a receiver function with the provided parsing_fn.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.data.build_ranking_serving_input_receiver_fn_with_parsing_fn" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py#L1005-L1032">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_sequence_example_serving_input_receiver_fn.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_sequence_example_serving_input_receiver_fn.md
index afe7376..aed6f81 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_sequence_example_serving_input_receiver_fn.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_sequence_example_serving_input_receiver_fn.md
@@ -1,3 +1,5 @@
+description: Creates a serving_input_receiver_fn for SequenceExample inputs.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.data.build_sequence_example_serving_input_receiver_fn" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py#L1248-L1279">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_tf_example_serving_input_receiver_fn.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_tf_example_serving_input_receiver_fn.md
new file mode 100644
index 0000000..ebac007
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/build_tf_example_serving_input_receiver_fn.md
@@ -0,0 +1,96 @@
+description: Builds a serving input fn for tensorflow.training.Example.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.data.build_tf_example_serving_input_receiver_fn" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# tfr.data.build_tf_example_serving_input_receiver_fn
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py#L1332-L1369">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Builds a serving input fn for `tensorflow.training.Example`.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.data.build_tf_example_serving_input_receiver_fn(
+    context_feature_spec, example_feature_spec, size_feature_name=None,
+    mask_feature_name=None, default_batch_size=None
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`context_feature_spec`
+</td>
+<td>
+(dict) Map from feature keys to `FixedLenFeature`,
+`VarLenFeature` or `RaggedFeature` values.
+</td>
+</tr><tr>
+<td>
+`example_feature_spec`
+</td>
+<td>
+(dict) Map from  feature keys to `FixedLenFeature`,
+`VarLenFeature` or `RaggedFeature` values.
+</td>
+</tr><tr>
+<td>
+`size_feature_name`
+</td>
+<td>
+(str) Name of feature for example list sizes. Populates
+the feature dictionary with a `tf.int32` Tensor of value 1, and of shape
+[batch_size] for this feature name. If None, which is default, this
+feature is not generated.
+</td>
+</tr><tr>
+<td>
+`mask_feature_name`
+</td>
+<td>
+(str) Name of feature for example list masks. Populates
+the feature dictionary with a `tf.bool` Tensor of shape [batch_size,
+list_size] for this feature name. If None, which is default, this feature
+is not generated.
+</td>
+</tr><tr>
+<td>
+`default_batch_size`
+</td>
+<td>
+(int) Number of instances expected per batch. Leave
+unset for variable batch size (recommended).
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
+<tr class="alt">
+<td colspan="2">
+A `tf.estimator.export.ServingInputReceiver` object, which packages the
+placeholders and the resulting feature Tensors together.
+</td>
+</tr>
+
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/make_parsing_fn.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/make_parsing_fn.md
index 360d544..e861059 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/make_parsing_fn.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/make_parsing_fn.md
@@ -1,3 +1,5 @@
+description: Returns a parsing fn for a standard data format.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.data.make_parsing_fn" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py#L800-L854">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -22,8 +23,8 @@ Returns a parsing fn for a standard data format.
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.data.make_parsing_fn(
     data_format, list_size=None, context_feature_spec=None,
-    example_feature_spec=None, size_feature_name=None, shuffle_examples=False,
-    seed=None
+    example_feature_spec=None, size_feature_name=None, mask_feature_name=None,
+    shuffle_examples=False, seed=None
 )
 </code></pre>
 
@@ -78,6 +79,16 @@ generated.
 </td>
 </tr><tr>
 <td>
+`mask_feature_name`
+</td>
+<td>
+(str) Name of feature for example list masks. Populates
+the feature dictionary with a `tf.bool` Tensor of shape [batch_size,
+list_size] for this feature name. If None, which is default, this feature
+is not generated.
+</td>
+</tr><tr>
+<td>
 `shuffle_examples`
 </td>
 <td>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_example_in_example.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_example_in_example.md
index 4e0a41a..f3cbf9b 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_example_in_example.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_example_in_example.md
@@ -1,3 +1,5 @@
+description: Parses an ExampleInExample batch to a feature map.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.data.parse_from_example_in_example" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py#L196-L365">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -22,8 +23,8 @@ Parses an ExampleInExample batch to a feature map.
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.data.parse_from_example_in_example(
     serialized, list_size=None, context_feature_spec=None,
-    example_feature_spec=None, size_feature_name=None, shuffle_examples=False,
-    seed=None
+    example_feature_spec=None, size_feature_name=None, mask_feature_name=None,
+    shuffle_examples=False, seed=None
 )
 </code></pre>
 
@@ -203,6 +204,16 @@ generated.
 </td>
 </tr><tr>
 <td>
+`mask_feature_name`
+</td>
+<td>
+(str) Name of feature for example list masks. Populates
+the feature dictionary with a `tf.bool` Tensor of shape [batch_size,
+list_size] for this feature name. If None, which is default, this feature
+is not generated.
+</td>
+</tr><tr>
+<td>
 `shuffle_examples`
 </td>
 <td>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_example_list.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_example_list.md
index 3800317..6a342d7 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_example_list.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_example_list.md
@@ -1,3 +1,5 @@
+description: Parses an ExampleListWithContext batch to a feature map.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.data.parse_from_example_list" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py#L376-L524">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -22,8 +23,8 @@ Parses an `ExampleListWithContext` batch to a feature map.
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.data.parse_from_example_list(
     serialized, list_size=None, context_feature_spec=None,
-    example_feature_spec=None, size_feature_name=None, shuffle_examples=False,
-    seed=None
+    example_feature_spec=None, size_feature_name=None, mask_feature_name=None,
+    shuffle_examples=False, seed=None
 )
 </code></pre>
 
@@ -183,6 +184,16 @@ generated.
 </td>
 </tr><tr>
 <td>
+`mask_feature_name`
+</td>
+<td>
+(str) Name of feature for example list masks. Populates
+the feature dictionary with a `tf.bool` Tensor of shape [batch_size,
+list_size] for this feature name. If None, which is default, this feature
+is not generated.
+</td>
+</tr><tr>
+<td>
 `shuffle_examples`
 </td>
 <td>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_sequence_example.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_sequence_example.md
index d8caee1..2e876ee 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_sequence_example.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_sequence_example.md
@@ -1,3 +1,5 @@
+description: Parses SequenceExample to feature maps.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.data.parse_from_sequence_example" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py#L656-L797">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -22,8 +23,8 @@ Parses SequenceExample to feature maps.
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.data.parse_from_sequence_example(
     serialized, list_size=None, context_feature_spec=None,
-    example_feature_spec=None, size_feature_name=None, shuffle_examples=False,
-    seed=None
+    example_feature_spec=None, size_feature_name=None, mask_feature_name=None,
+    shuffle_examples=False, seed=None
 )
 </code></pre>
 
@@ -177,6 +178,16 @@ generated.
 </td>
 </tr><tr>
 <td>
+`mask_feature_name`
+</td>
+<td>
+(str) Name of feature for example list masks. Populates
+the feature dictionary with a `tf.bool` Tensor of shape [batch_size,
+list_size] for this feature name. If None, which is default, this feature
+is not generated.
+</td>
+</tr><tr>
+<td>
 `shuffle_examples`
 </td>
 <td>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_tf_example.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_tf_example.md
new file mode 100644
index 0000000..b28417f
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/parse_from_tf_example.md
@@ -0,0 +1,97 @@
+description: Parse function to convert tf.train.Example to feature maps.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.data.parse_from_tf_example" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# tfr.data.parse_from_tf_example
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py#L1282-L1329">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Parse function to convert `tf.train.Example` to feature maps.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.data.parse_from_tf_example(
+    serialized, context_feature_spec=None, example_feature_spec=None,
+    size_feature_name=None, mask_feature_name=None
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`serialized`
+</td>
+<td>
+(`tf.train.Example`) A serialized proto object containing
+context and example features.
+</td>
+</tr><tr>
+<td>
+`context_feature_spec`
+</td>
+<td>
+(dict) A mapping from feature keys to
+`FixedLenFeature`, `VarLenFeature` or `RaggedFeature` values for context
+in `tf.train.Example` proto.
+</td>
+</tr><tr>
+<td>
+`example_feature_spec`
+</td>
+<td>
+(dict) A mapping from feature keys to
+`FixedLenFeature`, `VarLenFeature` or `RaggedFeature` values for examples
+in `tf.train.Example` proto.
+</td>
+</tr><tr>
+<td>
+`size_feature_name`
+</td>
+<td>
+(str) Name of feature for example list sizes. Populates
+the feature dictionary with a `tf.int32` Tensor of shape [batch_size] for
+this feature name. If None, which is default, this feature is not
+generated.
+</td>
+</tr><tr>
+<td>
+`mask_feature_name`
+</td>
+<td>
+(str) Name of feature for example list masks. Populates
+the feature dictionary with a `tf.bool` Tensor of shape [batch_size,
+list_size] for this feature name. If None, which is default, this feature
+is not generated.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
+<tr class="alt">
+<td colspan="2">
+A mapping from feature keys to `Tensor`, `SparseTensor` or `RaggedTensor`.
+</td>
+</tr>
+
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/read_batched_sequence_example_dataset.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/read_batched_sequence_example_dataset.md
index 56d6bb5..6b1a176 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/data/read_batched_sequence_example_dataset.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/data/read_batched_sequence_example_dataset.md
@@ -1,3 +1,5 @@
+description: Returns a Dataset of features from SequenceExample.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.data.read_batched_sequence_example_dataset" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/data.py#L1083-L1244">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator.md
index 25a7bf5..9770f40 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator.md
@@ -1,3 +1,6 @@
+description: Provides an EstimatorBuilder for creating a TF-Ranking model
+estimator.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.estimator" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +10,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/estimator.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/EstimatorBuilder.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/EstimatorBuilder.md
index 0e10a2c..d25352d 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/EstimatorBuilder.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/EstimatorBuilder.md
@@ -1,3 +1,5 @@
+description: Builds a tf.estimator.Estimator for a TF-Ranking model.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.estimator.EstimatorBuilder" />
 <meta itemprop="path" content="Stable" />
@@ -9,10 +11,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/estimator.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/estimator.py#L105-L365">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -201,7 +202,7 @@ are not specified.
 
 <h3 id="make_estimator"><code>make_estimator</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/estimator.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/estimator.py#L359-L365">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/GAMEstimatorBuilder.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/GAMEstimatorBuilder.md
index cea6583..e188088 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/GAMEstimatorBuilder.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/GAMEstimatorBuilder.md
@@ -1,3 +1,5 @@
+description: Builds a TFR estimator with subscore signatures of GAM models.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.estimator.GAMEstimatorBuilder" />
 <meta itemprop="path" content="Stable" />
@@ -9,10 +11,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/estimator.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/estimator.py#L670-L739">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -175,7 +176,7 @@ are not specified.
 
 <h3 id="make_estimator"><code>make_estimator</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/estimator.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/estimator.py#L359-L365">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/make_dnn_ranking_estimator.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/make_dnn_ranking_estimator.md
index bd650fe..ae7b4b3 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/make_dnn_ranking_estimator.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/make_dnn_ranking_estimator.md
@@ -1,3 +1,5 @@
+description: Builds an Estimator instance with DNN scoring function.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.estimator.make_dnn_ranking_estimator" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/estimator.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/estimator.py#L444-L516">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -22,12 +23,11 @@ Builds an `Estimator` instance with DNN scoring function.
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.estimator.make_dnn_ranking_estimator(
     example_feature_columns, hidden_units, context_feature_columns=None,
-    optimizer=None, learning_rate=0.05, listwise_inference=False,
-    loss='approx_ndcg_loss',
+    optimizer=None, learning_rate=0.05, loss=&#x27;approx_ndcg_loss&#x27;,
     loss_reduction=tf.compat.v1.losses.Reduction.SUM_OVER_BATCH_SIZE,
     activation_fn=tf.nn.relu, dropout=None, use_batch_norm=False,
     batch_norm_moment=0.999, model_dir=None, checkpoint_secs=120,
-    num_checkpoints=1000
+    num_checkpoints=1000, listwise_inference=False
 )
 </code></pre>
 
@@ -78,14 +78,6 @@ has 64 nodes and second one has 32.
 </td>
 </tr><tr>
 <td>
-`listwise_inference`
-</td>
-<td>
-(bool) Whether the inference will be performed with the
-listwise data format such as `ExampleListWithContext`.
-</td>
-</tr><tr>
-<td>
 `loss`
 </td>
 <td>
@@ -155,6 +147,14 @@ continue training a previously saved model.
 <td>
 (int) Number of checkpoints to keep.
 </td>
+</tr><tr>
+<td>
+`listwise_inference`
+</td>
+<td>
+(bool) Whether the inference will be performed with the
+listwise data format such as `ExampleListWithContext`.
+</td>
 </tr>
 </table>
 
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/make_gam_ranking_estimator.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/make_gam_ranking_estimator.md
index 905cf73..61dcfa1 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/make_gam_ranking_estimator.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/estimator/make_gam_ranking_estimator.md
@@ -1,3 +1,5 @@
+description: Builds an Estimator instance with GAM scoring function.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.estimator.make_gam_ranking_estimator" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/estimator.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/estimator.py#L742-L824">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -23,11 +24,11 @@ Builds an `Estimator` instance with GAM scoring function.
 <code>tfr.estimator.make_gam_ranking_estimator(
     example_feature_columns, example_hidden_units, context_feature_columns=None,
     context_hidden_units=None, optimizer=None, learning_rate=0.05,
-    loss='approx_ndcg_loss',
+    loss=&#x27;approx_ndcg_loss&#x27;,
     loss_reduction=tf.compat.v1.losses.Reduction.SUM_OVER_BATCH_SIZE,
     activation_fn=tf.nn.relu, dropout=None, use_batch_norm=False,
     batch_norm_moment=0.999, model_dir=None, checkpoint_secs=120,
-    num_checkpoints=1000
+    num_checkpoints=1000, listwise_inference=False
 )
 </code></pre>
 
@@ -161,6 +162,14 @@ continue training a previously saved model.
 <td>
 (int) Number of checkpoints to keep.
 </td>
+</tr><tr>
+<td>
+`listwise_inference`
+</td>
+<td>
+(bool) Whether the inference will be performed with the
+listwise data format such as `ExampleListWithContext`.
+</td>
 </tr>
 </table>
 
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/ext.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/ext.md
index 2a9af70..39e043f 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/ext.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/ext.md
@@ -1,3 +1,5 @@
+description: TensorFlow Ranking library.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.ext" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/__init__.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/pipeline.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/pipeline.md
index c690e58..9e0905b 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/pipeline.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/pipeline.md
@@ -1,3 +1,5 @@
+description: Provides a RankingPipeline for running a TF-Ranking model.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.ext.pipeline" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/pipeline.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/pipeline/RankingPipeline.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/pipeline/RankingPipeline.md
index 617b61d..0e4a197 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/pipeline/RankingPipeline.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/pipeline/RankingPipeline.md
@@ -1,3 +1,6 @@
+description: Class to set up the input, train and eval processes for a TF
+Ranking model.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.ext.pipeline.RankingPipeline" />
 <meta itemprop="path" content="Stable" />
@@ -9,10 +12,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/pipeline.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/pipeline.py#L20-L408">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -24,7 +26,7 @@ Class to set up the input, train and eval processes for a TF Ranking model.
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.ext.pipeline.RankingPipeline(
     context_feature_columns, example_feature_columns, hparams, estimator,
-    label_feature_name='relevance', label_feature_type=tf.int64,
+    label_feature_name=&#x27;relevance&#x27;, label_feature_type=tf.int64,
     dataset_reader=tf.data.TFRecordDataset, best_exporter_metric=None,
     best_exporter_metric_higher_better=True, size_feature_name=None
 )
@@ -218,7 +220,7 @@ generated.
 
 <h3 id="train_and_eval"><code>train_and_eval</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/pipeline.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/pipeline.py#L401-L408">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/tfrbert.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/tfrbert.md
index 78c9d18..2e4ad98 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/tfrbert.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/tfrbert.md
@@ -1,3 +1,6 @@
+description: Class defining utilities for finetuning TF-Ranking models with
+Bert.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.ext.tfrbert" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +10,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/tfrbert/TFRBertRankingNetwork.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/tfrbert/TFRBertRankingNetwork.md
index 9d4f40a..8df8e40 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/tfrbert/TFRBertRankingNetwork.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/tfrbert/TFRBertRankingNetwork.md
@@ -1,3 +1,5 @@
+description: A TFRBertRankingNetwork scoring based univariate ranking network.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.ext.tfrbert.TFRBertRankingNetwork" />
 <meta itemprop="path" content="Stable" />
@@ -24,10 +26,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py#L18-L112">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -37,12 +38,13 @@
 A TFRBertRankingNetwork scoring based univariate ranking network.
 
 Inherits From:
-[`UnivariateRankingNetwork`](../../../tfr/keras/network/UnivariateRankingNetwork.md)
+[`UnivariateRankingNetwork`](../../../tfr/keras/network/UnivariateRankingNetwork.md),
+[`RankingNetwork`](../../../tfr/keras/network/RankingNetwork.md)
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.ext.tfrbert.TFRBertRankingNetwork(
     context_feature_columns, example_feature_columns, bert_config_file,
-    bert_max_seq_length, bert_output_dropout, name='tfrbert', **kwargs
+    bert_max_seq_length, bert_output_dropout, name=&#x27;tfrbert&#x27;, **kwargs
 )
 </code></pre>
 
@@ -118,13 +120,32 @@ keyword arguments.
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `context_feature_columns` </td>
-<td>
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`context_feature_columns` </td> <td>
+
+</td> </tr><tr> <td> `dtype` </td> <td> The dtype of the layer weights.
 
-</td> </tr><tr> <td> `dtype` </td> <td> Dtype used by the weights of the layer,
-set in the constructor. </td> </tr><tr> <td> `dynamic` </td> <td> Whether the
-layer is dynamic (eager-only); set in the constructor. </td> </tr><tr> <td>
-`example_feature_columns` </td> <td>
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `example_feature_columns` </td> <td>
 
 </td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
 layer.
@@ -251,6 +272,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -322,9 +350,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -363,11 +392,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -446,9 +475,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -464,7 +494,7 @@ Instance of `TensorShape`, or list of instances of
 
 <h3 id="compute_logits"><code>compute_logits</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L252-L289">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -553,9 +583,10 @@ If `scorer` does not return a scalar output.
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -602,9 +633,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -669,7 +701,7 @@ if the layer isn't yet built
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L198-L216">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -718,7 +750,7 @@ A RankingNetwork layer.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py#L104-L112">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -734,6 +766,10 @@ its trained weights) from this configuration.
 The config of a layer does not include connectivity information, nor the layer
 class name. These are handled by `Network` (one layer of abstraction above).
 
+Note that `get_config()` does not guarantee to return a fresh copy of dict every
+time it is called. The callers should make a copy of the returned dict if they
+want to modify it.
+
 <!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
@@ -798,7 +834,7 @@ Weights values as a list of numpy arrays.
 
 <h3 id="score"><code>score</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py#L78-L102">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -893,9 +929,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -929,7 +966,7 @@ layer's specifications.
 
 <h3 id="transform"><code>transform</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L118-L141">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -1068,9 +1105,10 @@ The original method wrapped such that it enters the module's name scope.
 Wraps `call`, applying pre- and post-processing steps.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -1111,6 +1149,7 @@ Output tensor(s).
     do), its default value will be set to the mask generated for `inputs` by the
     previous layer (if `input` did come from a layer that generated a
     corresponding mask, i.e. if it came from a Keras layer with masking support.
+-   If the layer is not built, the method will call `build`.
 
 <!-- Tabular view -->
  <table class="responsive fixed orange">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/tfrbert/TFRBertUtil.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/tfrbert/TFRBertUtil.md
index 5b8fbc9..8b9e23b 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/tfrbert/TFRBertUtil.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/ext/tfrbert/TFRBertUtil.md
@@ -1,3 +1,5 @@
+description: Class that defines a set of utility functions for Bert.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.ext.tfrbert.TFRBertUtil" />
 <meta itemprop="path" content="Stable" />
@@ -11,10 +13,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py#L115-L331">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -70,7 +71,7 @@ This should be aligned with the `vocab_file`.
 
 <h3 id="convert_to_elwc"><code>convert_to_elwc</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py#L297-L331">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -132,12 +133,12 @@ A tensorflow.serving.ExampleListWithContext example containing the
 
 <h3 id="create_optimizer"><code>create_optimizer</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py#L145-L170">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>create_optimizer(
-    init_lr, train_steps, warmup_steps, optimizer_type='adamw'
+    init_lr, train_steps, warmup_steps, optimizer_type=&#x27;adamw&#x27;
 )
 </code></pre>
 
@@ -197,7 +198,7 @@ The optimizer training op.
 
 <h3 id="get_warm_start_settings"><code>get_warm_start_settings</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/extension/tfrbert.py#L172-L194">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/feature.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/feature.md
index e92dbcf..3e58d86 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/feature.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/feature.md
@@ -1,3 +1,5 @@
+description: Feature transformations for ranking library.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.feature" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/feature.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/encode_features.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/encode_features.md
index 84242ef..ce0e5f3 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/encode_features.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/encode_features.md
@@ -1,3 +1,5 @@
+description: Returns dense tensors from features using feature columns.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.feature.encode_features" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/feature.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/feature.py#L64-L118">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/encode_listwise_features.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/encode_listwise_features.md
index 148a3f5..430eb4f 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/encode_listwise_features.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/encode_listwise_features.md
@@ -1,3 +1,5 @@
+description: Returns dense tensors from features using feature columns.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.feature.encode_listwise_features" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/feature.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/feature.py#L121-L207">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/encode_pointwise_features.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/encode_pointwise_features.md
index 1d5a486..b59f01f 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/encode_pointwise_features.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/encode_pointwise_features.md
@@ -1,3 +1,6 @@
+description: Returns dense tensors from pointwise features using feature
+columns.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.feature.encode_pointwise_features" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +10,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/feature.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/feature.py#L210-L251">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/make_identity_transform_fn.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/make_identity_transform_fn.md
index 4da1843..7bbb3fe 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/make_identity_transform_fn.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/feature/make_identity_transform_fn.md
@@ -1,3 +1,5 @@
+description: Returns transform fn that split the features.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.feature.make_identity_transform_fn" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/feature.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/feature.py#L14-L48">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/head.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/head.md
index c1682a3..ba6ac83 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/head.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/head.md
@@ -1,3 +1,5 @@
+description: Defines Heads of TF ranking models.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.head" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/head.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/head/create_multi_ranking_head.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/head/create_multi_ranking_head.md
index 87302df..3011c1c 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/head/create_multi_ranking_head.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/head/create_multi_ranking_head.md
@@ -1,3 +1,5 @@
+description: A factory method to create _MultiRankingHead.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.head.create_multi_ranking_head" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/head.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/head.py#L75-L85">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/head/create_ranking_head.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/head/create_ranking_head.md
index a981379..3cbc706 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/head/create_ranking_head.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/head/create_ranking_head.md
@@ -1,3 +1,5 @@
+description: A factory method to create _RankingHead.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.head.create_ranking_head" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/head.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/head.py#L23-L72">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras.md
index 9a5e45f..134553c 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras.md
@@ -1,3 +1,5 @@
+description: TensorFlow Ranking Keras Module.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/__init__.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
@@ -30,6 +31,8 @@ and estimator.
 [`feature`](../tfr/keras/feature.md) module: Feature transformations for ranking
 in Keras.
 
+[`layers`](../tfr/keras/layers.md) module: Defines Keras Layers for TF-Ranking.
+
 [`losses`](../tfr/keras/losses.md) module: Keras losses in TF-Ranking.
 
 [`metrics`](../tfr/keras/metrics.md) module: Keras metrics in TF-Ranking.
@@ -38,3 +41,8 @@ in Keras.
 Keras.
 
 [`network`](../tfr/keras/network.md) module: Ranking Networks in Keras.
+
+[`saved_model`](../tfr/keras/saved_model.md) module: SavedModel utilities for
+TF-Ranking.
+
+[`utils`](../tfr/keras/utils.md) module: Utils for tfr.keras.
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned.md
index 6db5992..0785e34 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned.md
@@ -1,3 +1,5 @@
+description: TensorFlow Ranking Premade Keras Module.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.canned" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/__init__.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/DNNRankingNetwork.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/DNNRankingNetwork.md
index 3b16737..7384af8 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/DNNRankingNetwork.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/DNNRankingNetwork.md
@@ -1,3 +1,5 @@
+description: Deep Neural Network (DNN) scoring based univariate ranking network.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.canned.DNNRankingNetwork" />
 <meta itemprop="path" content="Stable" />
@@ -24,10 +26,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/dnn.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/dnn.py#L9-L123">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -37,7 +38,8 @@
 Deep Neural Network (DNN) scoring based univariate ranking network.
 
 Inherits From:
-[`UnivariateRankingNetwork`](../../../tfr/keras/network/UnivariateRankingNetwork.md)
+[`UnivariateRankingNetwork`](../../../tfr/keras/network/UnivariateRankingNetwork.md),
+[`RankingNetwork`](../../../tfr/keras/network/RankingNetwork.md)
 
 <section class="expandable">
   <h4 class="showalways">View aliases</h4>
@@ -51,7 +53,8 @@ Inherits From:
 <code>tfr.keras.canned.DNNRankingNetwork(
     context_feature_columns=None, example_feature_columns=None,
     hidden_layer_dims=None, activation=None, use_batch_norm=True,
-    batch_norm_moment=0.999, dropout=0.5, name='dnn_ranking_network', **kwargs
+    batch_norm_moment=0.999, dropout=0.5, name=&#x27;dnn_ranking_network&#x27;,
+    **kwargs
 )
 </code></pre>
 
@@ -155,13 +158,32 @@ keyword arguments.
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `context_feature_columns` </td>
-<td>
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`context_feature_columns` </td> <td>
+
+</td> </tr><tr> <td> `dtype` </td> <td> The dtype of the layer weights.
 
-</td> </tr><tr> <td> `dtype` </td> <td> Dtype used by the weights of the layer,
-set in the constructor. </td> </tr><tr> <td> `dynamic` </td> <td> Whether the
-layer is dynamic (eager-only); set in the constructor. </td> </tr><tr> <td>
-`example_feature_columns` </td> <td>
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `example_feature_columns` </td> <td>
 
 </td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
 layer.
@@ -288,6 +310,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -359,9 +388,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -400,11 +430,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -483,9 +513,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -501,7 +532,7 @@ Instance of `TensorShape`, or list of instances of
 
 <h3 id="compute_logits"><code>compute_logits</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L252-L289">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -590,9 +621,10 @@ If `scorer` does not return a scalar output.
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -639,9 +671,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -706,7 +739,7 @@ if the layer isn't yet built
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L198-L216">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -755,7 +788,7 @@ A RankingNetwork layer.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/dnn.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/dnn.py#L114-L123">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -771,6 +804,10 @@ its trained weights) from this configuration.
 The config of a layer does not include connectivity information, nor the layer
 class name. These are handled by `Network` (one layer of abstraction above).
 
+Note that `get_config()` does not guarantee to return a fresh copy of dict every
+time it is called. The callers should make a copy of the returned dict if they
+want to modify it.
+
 <!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
@@ -835,7 +872,7 @@ Weights values as a list of numpy arrays.
 
 <h3 id="score"><code>score</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/dnn.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/dnn.py#L84-L112">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -930,9 +967,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -966,7 +1004,7 @@ layer's specifications.
 
 <h3 id="transform"><code>transform</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L118-L141">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -1105,9 +1143,10 @@ The original method wrapped such that it enters the module's name scope.
 Wraps `call`, applying pre- and post-processing steps.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -1148,6 +1187,7 @@ Output tensor(s).
     do), its default value will be set to the mask generated for `inputs` by the
     previous layer (if `input` did come from a layer that generated a
     corresponding mask, i.e. if it came from a Keras layer with masking support.
+-   If the layer is not built, the method will call `build`.
 
 <!-- Tabular view -->
  <table class="responsive fixed orange">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/GAMRankingNetwork.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/GAMRankingNetwork.md
index 6c2f575..c0e7289 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/GAMRankingNetwork.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/GAMRankingNetwork.md
@@ -1,3 +1,5 @@
+description: Generalized Additive Model (GAM) based univariate ranking network.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.canned.GAMRankingNetwork" />
 <meta itemprop="path" content="Stable" />
@@ -24,10 +26,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/gam.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/gam.py#L56-L217">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -37,7 +38,8 @@
 Generalized Additive Model (GAM) based univariate ranking network.
 
 Inherits From:
-[`UnivariateRankingNetwork`](../../../tfr/keras/network/UnivariateRankingNetwork.md)
+[`UnivariateRankingNetwork`](../../../tfr/keras/network/UnivariateRankingNetwork.md),
+[`RankingNetwork`](../../../tfr/keras/network/RankingNetwork.md)
 
 <section class="expandable">
   <h4 class="showalways">View aliases</h4>
@@ -52,7 +54,7 @@ Inherits From:
     context_feature_columns=None, example_feature_columns=None,
     context_hidden_layer_dims=None, example_hidden_layer_dims=None, activation=None,
     use_batch_norm=True, batch_norm_moment=0.999, dropout=0.5,
-    name='gam_ranking_model', **kwargs
+    name=&#x27;gam_ranking_model&#x27;, **kwargs
 )
 </code></pre>
 
@@ -165,13 +167,32 @@ Keyword arguments.
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `context_feature_columns` </td>
-<td>
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`context_feature_columns` </td> <td>
+
+</td> </tr><tr> <td> `dtype` </td> <td> The dtype of the layer weights.
 
-</td> </tr><tr> <td> `dtype` </td> <td> Dtype used by the weights of the layer,
-set in the constructor. </td> </tr><tr> <td> `dynamic` </td> <td> Whether the
-layer is dynamic (eager-only); set in the constructor. </td> </tr><tr> <td>
-`example_feature_columns` </td> <td>
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `example_feature_columns` </td> <td>
 
 </td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
 layer.
@@ -298,6 +319,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -369,9 +397,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -410,11 +439,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -493,9 +522,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -511,7 +541,7 @@ Instance of `TensorShape`, or list of instances of
 
 <h3 id="compute_logits"><code>compute_logits</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L252-L289">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -600,9 +630,10 @@ If `scorer` does not return a scalar output.
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -649,9 +680,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -716,7 +748,7 @@ if the layer isn't yet built
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L198-L216">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -765,7 +797,7 @@ A RankingNetwork layer.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/gam.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/gam.py#L207-L217">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -781,6 +813,10 @@ its trained weights) from this configuration.
 The config of a layer does not include connectivity information, nor the layer
 class name. These are handled by `Network` (one layer of abstraction above).
 
+Note that `get_config()` does not guarantee to return a fresh copy of dict every
+time it is called. The callers should make a copy of the returned dict if they
+want to modify it.
+
 <!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
@@ -845,7 +881,7 @@ Weights values as a list of numpy arrays.
 
 <h3 id="score"><code>score</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/gam.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/gam.py#L140-L205">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -940,9 +976,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -976,7 +1013,7 @@ layer's specifications.
 
 <h3 id="transform"><code>transform</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L118-L141">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -1115,9 +1152,10 @@ The original method wrapped such that it enters the module's name scope.
 Wraps `call`, applying pre- and post-processing steps.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -1158,6 +1196,7 @@ Output tensor(s).
     do), its default value will be set to the mask generated for `inputs` by the
     previous layer (if `input` did come from a layer that generated a
     corresponding mask, i.e. if it came from a Keras layer with masking support.
+-   If the layer is not built, the method will call `build`.
 
 <!-- Tabular view -->
  <table class="responsive fixed orange">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/dnn.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/dnn.md
index 2a2903d..61a1a54 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/dnn.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/dnn.md
@@ -1,3 +1,5 @@
+description: DNN Ranking network in Keras.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.canned.dnn" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/dnn.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/gam.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/gam.md
index bfab34d..591dd35 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/gam.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/canned/gam.md
@@ -1,3 +1,5 @@
+description: Generalized Additive Model (GAM) Keras Network.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.canned.gam" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/canned/gam.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/estimator.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/estimator.md
index 5d3c2aa..f7ee503 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/estimator.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/estimator.md
@@ -1,3 +1,5 @@
+description: Adaptor between keras models and estimator.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.estimator" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/estimator.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/estimator/model_to_estimator.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/estimator/model_to_estimator.md
index adf78e2..b4b4409 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/estimator/model_to_estimator.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/estimator/model_to_estimator.md
@@ -1,3 +1,5 @@
+description: Keras ranking model to Estimator.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.estimator.model_to_estimator" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/estimator.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/estimator.py#L17-L161">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -97,7 +98,6 @@ mode.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature.md
index d12f19b..7a1def0 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature.md
@@ -1,3 +1,5 @@
+description: Feature transformations for ranking in Keras.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.feature" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/EncodeListwiseFeatures.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/EncodeListwiseFeatures.md
index b054293..e50c487 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/EncodeListwiseFeatures.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/EncodeListwiseFeatures.md
@@ -1,3 +1,6 @@
+description: A layer that produces dense Tensors from context and example
+features.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.feature.EncodeListwiseFeatures" />
 <meta itemprop="path" content="Stable" />
@@ -21,10 +24,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py#L184-L306">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -36,7 +38,7 @@ A layer that produces dense `Tensors` from context and example features.
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.feature.EncodeListwiseFeatures(
     context_feature_columns=None, example_feature_columns=None,
-    name='listwise_dense_features', **kwargs
+    name=&#x27;listwise_dense_features&#x27;, **kwargs
 )
 </code></pre>
 
@@ -84,13 +86,32 @@ keyword arguments.
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `context_feature_columns` </td>
-<td>
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`context_feature_columns` </td> <td>
+
+</td> </tr><tr> <td> `dtype` </td> <td> The dtype of the layer weights.
 
-</td> </tr><tr> <td> `dtype` </td> <td> Dtype used by the weights of the layer,
-set in the constructor. </td> </tr><tr> <td> `dynamic` </td> <td> Whether the
-layer is dynamic (eager-only); set in the constructor. </td> </tr><tr> <td>
-`example_feature_columns` </td> <td>
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `example_feature_columns` </td> <td>
 
 </td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
 layer.
@@ -217,6 +238,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -288,9 +316,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -329,11 +358,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -412,9 +441,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -439,9 +469,10 @@ Instance of `TensorShape`, or list of instances of
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -488,9 +519,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -555,7 +587,7 @@ if the layer isn't yet built
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py#L288-L306">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -604,7 +636,7 @@ A EncodeListwiseFeatures layer.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py#L277-L286">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -620,6 +652,10 @@ its trained weights) from this configuration.
 The config of a layer does not include connectivity information, nor the layer
 class name. These are handled by `Network` (one layer of abstraction above).
 
+Note that `get_config()` does not guarantee to return a fresh copy of dict every
+time it is called. The callers should make a copy of the returned dict if they
+want to modify it.
+
 <!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
@@ -723,9 +759,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -742,7 +779,6 @@ output of `get_weights`).
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Raises</th></tr>
@@ -828,9 +864,10 @@ The original method wrapped such that it enters the module's name scope.
 Wraps `call`, applying pre- and post-processing steps.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -871,6 +908,7 @@ Output tensor(s).
     do), its default value will be set to the mask generated for `inputs` by the
     previous layer (if `input` did come from a layer that generated a
     corresponding mask, i.e. if it came from a Keras layer with masking support.
+-   If the layer is not built, the method will call `build`.
 
 <!-- Tabular view -->
  <table class="responsive fixed orange">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/GenerateMask.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/GenerateMask.md
index d59d5e1..8fb44c2 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/GenerateMask.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/GenerateMask.md
@@ -1,3 +1,5 @@
+description: Layer to generate mask.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.feature.GenerateMask" />
 <meta itemprop="path" content="Stable" />
@@ -21,10 +23,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py#L108-L181">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -35,7 +36,8 @@ Layer to generate mask.
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.feature.GenerateMask(
-    example_feature_columns, size_feature_name, name='generate_mask_layer', **kwargs
+    example_feature_columns, size_feature_name,
+    name=&#x27;generate_mask_layer&#x27;, **kwargs
 )
 </code></pre>
 
@@ -86,10 +88,30 @@ keyword arguments.
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `dtype` </td> <td> Dtype used by
-the weights of the layer, set in the constructor. </td> </tr><tr> <td> `dynamic`
-</td> <td> Whether the layer is dynamic (eager-only); set in the constructor.
-</td> </tr><tr> <td> `example_feature_columns` </td> <td>
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `example_feature_columns` </td> <td>
 
 </td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
 layer.
@@ -217,6 +239,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -288,9 +317,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -329,11 +359,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -412,9 +442,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -439,9 +470,10 @@ Instance of `TensorShape`, or list of instances of
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -488,9 +520,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -555,7 +588,7 @@ if the layer isn't yet built
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py#L166-L181">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -604,7 +637,7 @@ A GenerateMask layer.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py#L156-L164">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -620,6 +653,10 @@ its trained weights) from this configuration.
 The config of a layer does not include connectivity information, nor the layer
 class name. These are handled by `Network` (one layer of abstraction above).
 
+Note that `get_config()` does not guarantee to return a fresh copy of dict every
+time it is called. The callers should make a copy of the returned dict if they
+want to modify it.
+
 <!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
@@ -723,9 +760,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -742,7 +780,6 @@ output of `get_weights`).
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Raises</th></tr>
@@ -828,9 +865,10 @@ The original method wrapped such that it enters the module's name scope.
 Wraps `call`, applying pre- and post-processing steps.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -871,6 +909,7 @@ Output tensor(s).
     do), its default value will be set to the mask generated for `inputs` by the
     previous layer (if `input` did come from a layer that generated a
     corresponding mask, i.e. if it came from a Keras layer with masking support.
+-   If the layer is not built, the method will call `build`.
 
 <!-- Tabular view -->
  <table class="responsive fixed orange">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/create_keras_inputs.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/create_keras_inputs.md
index d51d6aa..4dbeb6e 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/create_keras_inputs.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/create_keras_inputs.md
@@ -1,3 +1,5 @@
+description: Create Keras inputs from context and example feature columns.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.feature.create_keras_inputs" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py#L63-L105">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/deserialize_feature_columns.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/deserialize_feature_columns.md
index 2653016..cdd3c96 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/deserialize_feature_columns.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/deserialize_feature_columns.md
@@ -1,3 +1,5 @@
+description: Deserializes dict of feature column configs.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.feature.deserialize_feature_columns" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py#L37-L60">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/serialize_feature_columns.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/serialize_feature_columns.md
index e83e410..6852892 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/serialize_feature_columns.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/feature/serialize_feature_columns.md
@@ -1,3 +1,5 @@
+description: Serializes feature columns to a dict of class name and config.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.feature.serialize_feature_columns" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/feature.py#L11-L34">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers.md
new file mode 100644
index 0000000..6ba8b76
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers.md
@@ -0,0 +1,35 @@
+description: Defines Keras Layers for TF-Ranking.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.layers" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# Module: tfr.keras.layers
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Defines Keras Layers for TF-Ranking.
+
+## Classes
+
+[`class ConcatFeatures`](../../tfr/keras/layers/ConcatFeatures.md): Concatenates
+context features and example features in a listwise manner.
+
+[`class DocumentInteractionAttention`](../../tfr/keras/layers/DocumentInteractionAttention.md):
+Cross Document Interaction Attention layer.
+
+[`class FlattenList`](../../tfr/keras/layers/FlattenList.md): Layer to flatten
+the example list.
+
+[`class RestoreList`](../../tfr/keras/layers/RestoreList.md): Output layer to
+restore listwise output shape.
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/ConcatFeatures.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/ConcatFeatures.md
new file mode 100644
index 0000000..b99a87e
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/ConcatFeatures.md
@@ -0,0 +1,960 @@
+description: Concatenates context features and example features in a listwise
+manner.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.layers.ConcatFeatures" />
+<meta itemprop="path" content="Stable" />
+<meta itemprop="property" content="__call__"/>
+<meta itemprop="property" content="__init__"/>
+<meta itemprop="property" content="__new__"/>
+<meta itemprop="property" content="add_loss"/>
+<meta itemprop="property" content="add_metric"/>
+<meta itemprop="property" content="build"/>
+<meta itemprop="property" content="compute_mask"/>
+<meta itemprop="property" content="compute_output_shape"/>
+<meta itemprop="property" content="count_params"/>
+<meta itemprop="property" content="from_config"/>
+<meta itemprop="property" content="get_config"/>
+<meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="set_weights"/>
+<meta itemprop="property" content="with_name_scope"/>
+</div>
+
+# tfr.keras.layers.ConcatFeatures
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L203-L293">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Concatenates context features and example features in a listwise manner.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.layers.ConcatFeatures(
+    circular_padding: bool = True,
+    name: Optional[str] = None,
+    **kwargs
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+
+Given dicts of `context features`, `example features`, this layer expands
+list_size times for the `context_features` and concatenates them with example
+features along the `list_size` axis. The output is a 3-d tensor with shape
+[batch_size, list_size, sum(feature_dims)], where sum(feature_dims) is the sum
+of all example feature dimensions and the context feature dimension.
+
+#### Example:
+
+```python
+  # Batch size = 2, list_size = 2.
+  context_features = {
+      'context_feature_1': [[1.], [2.]]
+  }
+  example_features = {
+      'example_feature_1':
+          [[[1., 0.], [0., 1.]], [[0., 1.], [1., 0.]]]
+  }
+  mask = [[True, False], [True, True]]
+  ConcatFeatures()(context_features, example_features, mask)
+  # Returns: [[[1., 1., 0.], [1., 1., 0.]], [[2., 0., 1.], [2., 1., 0.]]])
+
+  ConcatFeatures(circular_padding=False)(
+      context_features, example_features, mask)
+  # Returns: [[[1., 1., 0.], [1., 0., 1.]], [[2., 0., 1.], [2., 1., 0.]]]
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`circular_padding`
+</td>
+<td>
+Whether to apply circular padding to replace invalid
+features with valid ones.
+</td>
+</tr><tr>
+<td>
+`name`
+</td>
+<td>
+Name of the layer.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+keyword arguments.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
+
+<tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `input` </td> <td> Retrieves the input
+tensor(s) of a layer.
+
+Only applicable if the layer has exactly one input, i.e. if it is connected to
+one incoming layer. </td> </tr><tr> <td> `input_spec` </td> <td> `InputSpec`
+instance(s) describing the input format for this layer.
+
+When you create a layer subclass, you can set `self.input_spec` to enable the
+layer to run input compatibility checks when it is called. Consider a `Conv2D`
+layer: it can only be called on a single input tensor of rank 4. As such, you
+can set, in `__init__()`:
+
+```python
+self.input_spec = tf.keras.layers.InputSpec(ndim=4)
+```
+
+Now, if you try to call the layer on an input that isn't rank 4 (for instance,
+an input of shape `(2,)`, it will raise a nicely-formatted error:
+
+```
+ValueError: Input 0 of layer conv2d is incompatible with the layer:
+expected ndim=4, found ndim=1. Full shape received: [2]
+```
+
+Input checks that can be specified via `input_spec` include: - Structure (e.g. a
+single input, a list of 2 inputs, etc) - Shape - Rank (ndim) - Dtype
+
+For more information, see `tf.keras.layers.InputSpec`. </td> </tr><tr> <td>
+`losses` </td> <td> List of losses added using the `add_loss()` API.
+
+Variable regularization tensors are created when this property is accessed, so
+it is eager safe: accessing `losses` under a `tf.GradientTape` will propagate
+gradients back to the corresponding variables.
+
+```
+>>> class MyLayer(tf.keras.layers.Layer):
+...   def call(self, inputs):
+...     self.add_loss(tf.abs(tf.reduce_mean(inputs)))
+...     return inputs
+>>> l = MyLayer()
+>>> l(np.ones((10, 1)))
+>>> l.losses
+[1.0]
+```
+
+```
+>>> inputs = tf.keras.Input(shape=(10,))
+>>> x = tf.keras.layers.Dense(10)(inputs)
+>>> outputs = tf.keras.layers.Dense(1)(x)
+>>> model = tf.keras.Model(inputs, outputs)
+>>> # Activity regularization.
+>>> len(model.losses)
+0
+>>> model.add_loss(tf.abs(tf.reduce_mean(x)))
+>>> len(model.losses)
+1
+```
+
+```
+>>> inputs = tf.keras.Input(shape=(10,))
+>>> d = tf.keras.layers.Dense(10, kernel_initializer='ones')
+>>> x = d(inputs)
+>>> outputs = tf.keras.layers.Dense(1)(x)
+>>> model = tf.keras.Model(inputs, outputs)
+>>> # Weight regularization.
+>>> model.add_loss(lambda: tf.reduce_mean(d.kernel))
+>>> model.losses
+[<tf.Tensor: shape=(), dtype=float32, numpy=1.0>]
+```
+
+</td> </tr><tr> <td> `metrics` </td> <td> List of metrics added using the
+`add_metric()` API.
+
+```
+>>> input = tf.keras.layers.Input(shape=(3,))
+>>> d = tf.keras.layers.Dense(2)
+>>> output = d(input)
+>>> d.add_metric(tf.reduce_max(output), name='max')
+>>> d.add_metric(tf.reduce_min(output), name='min')
+>>> [m.name for m in d.metrics]
+['max', 'min']
+```
+
+</td> </tr><tr> <td> `name` </td> <td> Name of the layer (string), set in the
+constructor. </td> </tr><tr> <td> `name_scope` </td> <td> Returns a
+`tf.name_scope` instance for this class. </td> </tr><tr> <td>
+`non_trainable_weights` </td> <td> List of all non-trainable weights tracked by
+this layer.
+
+Non-trainable weights are *not* updated during training. They are expected to be
+updated manually in `call()`. </td> </tr><tr> <td> `output` </td> <td> Retrieves
+the output tensor(s) of a layer.
+
+Only applicable if the layer has exactly one output, i.e. if it is connected to
+one incoming layer. </td> </tr><tr> <td> `submodules` </td> <td> Sequence of all
+sub-modules.
+
+Submodules are modules which are properties of this module, or found as
+properties of modules which are properties of this module (and so on).
+
+```
+>>> a = tf.Module()
+>>> b = tf.Module()
+>>> c = tf.Module()
+>>> a.b = b
+>>> b.c = c
+>>> list(a.submodules) == [b, c]
+True
+>>> list(b.submodules) == [c]
+True
+>>> list(c.submodules) == []
+True
+```
+
+</td> </tr><tr> <td> `supports_masking` </td> <td> Whether this layer supports
+computing a mask using `compute_mask`. </td> </tr><tr> <td> `trainable` </td>
+<td>
+
+</td> </tr><tr> <td> `trainable_weights` </td> <td> List of all trainable
+weights tracked by this layer.
+
+Trainable weights are updated via gradient descent during training.
+</td>
+</tr><tr>
+<td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
+`weights`
+</td>
+<td>
+Returns the list of all layer variables/weights.
+</td>
+</tr>
+</table>
+
+## Methods
+
+<h3 id="add_loss"><code>add_loss</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>add_loss(
+    losses, **kwargs
+)
+</code></pre>
+
+Add loss tensor(s), potentially dependent on layer inputs.
+
+Some losses (for instance, activity regularization losses) may be dependent on
+the inputs passed when calling a layer. Hence, when reusing the same layer on
+different inputs `a` and `b`, some entries in `layer.losses` may be dependent on
+`a` and some on `b`. This method automatically keeps track of dependencies.
+
+This method can be used inside a subclassed layer or model's `call` function, in
+which case `losses` should be a Tensor or list of Tensors.
+
+#### Example:
+
+```python
+class MyLayer(tf.keras.layers.Layer):
+  def call(self, inputs):
+    self.add_loss(tf.abs(tf.reduce_mean(inputs)))
+    return inputs
+```
+
+This method can also be called directly on a Functional Model during
+construction. In this case, any loss Tensors passed to this Model must be
+symbolic and be able to be traced back to the model's `Input`s. These losses
+become part of the model's topology and are tracked in `get_config`.
+
+#### Example:
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+# Activity regularization.
+model.add_loss(tf.abs(tf.reduce_mean(x)))
+```
+
+If this is not the case for your loss (if, for example, your loss references a
+`Variable` of one of the model's layers), you can wrap your loss in a
+zero-argument lambda. These losses are not tracked as part of the model's
+topology since they can't be serialized.
+
+#### Example:
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+d = tf.keras.layers.Dense(10)
+x = d(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+# Weight regularization.
+model.add_loss(lambda: tf.reduce_mean(d.kernel))
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`losses`
+</td>
+<td>
+Loss tensor, or list/tuple of tensors. Rather than tensors, losses
+may also be zero-argument callables which create a loss tensor.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Additional keyword arguments for backward compatibility.
+Accepted values:
+inputs - Deprecated, will be automatically inferred.
+</td>
+</tr>
+</table>
+
+<h3 id="add_metric"><code>add_metric</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>add_metric(
+    value, name=None, **kwargs
+)
+</code></pre>
+
+Adds metric tensor to the layer.
+
+This method can be used inside the `call()` method of a subclassed layer or
+model.
+
+```python
+class MyMetricLayer(tf.keras.layers.Layer):
+  def __init__(self):
+    super(MyMetricLayer, self).__init__(name='my_metric_layer')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
+
+  def call(self, inputs):
+    self.add_metric(self.mean(x))
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
+    return inputs
+```
+
+This method can also be called directly on a Functional Model during
+construction. In this case, any tensor passed to this Model must be symbolic and
+be able to be traced back to the model's `Input`s. These metrics become part of
+the model's topology and are tracked when you save the model via `save()`.
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+model.add_metric(math_ops.reduce_sum(x), name='metric_1')
+```
+
+Note: Calling `add_metric()` with the result of a metric object on a Functional
+Model, as shown in the example below, is not supported. This is because we
+cannot trace the metric result tensor back to the model's inputs.
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+model.add_metric(tf.keras.metrics.Mean()(x), name='metric_1')
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`value`
+</td>
+<td>
+Metric tensor.
+</td>
+</tr><tr>
+<td>
+`name`
+</td>
+<td>
+String metric name.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Additional keyword arguments for backward compatibility.
+Accepted values:
+`aggregation` - When the `value` tensor provided is not the result of
+calling a `keras.Metric` instance, it will be aggregated by default
+using a `keras.Metric.Mean`.
+</td>
+</tr>
+</table>
+
+<h3 id="build"><code>build</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>build(
+    input_shape
+)
+</code></pre>
+
+Creates the variables of the layer (optional, for subclass implementers).
+
+This is a method that implementers of subclasses of `Layer` or `Model` can
+override if they need a state-creation step in-between layer instantiation and
+layer call.
+
+This is typically used to create the weights of `Layer` subclasses.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`input_shape`
+</td>
+<td>
+Instance of `TensorShape`, or list of instances of
+`TensorShape` if the layer expects a list of inputs
+(one instance per input).
+</td>
+</tr>
+</table>
+
+<h3 id="compute_mask"><code>compute_mask</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>compute_mask(
+    inputs, mask=None
+)
+</code></pre>
+
+Computes an output mask tensor.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`inputs`
+</td>
+<td>
+Tensor or list of tensors.
+</td>
+</tr><tr>
+<td>
+`mask`
+</td>
+<td>
+Tensor or list of tensors.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+None or a tensor (or list of tensors,
+one per output tensor of the layer).
+</td>
+</tr>
+
+</table>
+
+<h3 id="compute_output_shape"><code>compute_output_shape</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>compute_output_shape(
+    input_shape
+)
+</code></pre>
+
+Computes the output shape of the layer.
+
+If the layer has not been built, this method will call `build` on the layer.
+This assumes that the layer will later be used with inputs that match the input
+shape provided here.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`input_shape`
+</td>
+<td>
+Shape tuple (tuple of integers)
+or list of shape tuples (one per output tensor of the layer).
+Shape tuples can include None for free dimensions,
+instead of an integer.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+An input shape tuple.
+</td>
+</tr>
+
+</table>
+
+<h3 id="count_params"><code>count_params</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>count_params()
+</code></pre>
+
+Count the total number of scalars composing the weights.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+An integer count.
+</td>
+</tr>
+
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Raises</th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+if the layer isn't yet built
+(in which case its weights aren't yet defined).
+</td>
+</tr>
+</table>
+
+<h3 id="from_config"><code>from_config</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@classmethod</code>
+<code>from_config(
+    config
+)
+</code></pre>
+
+Creates a layer from its config.
+
+This method is the reverse of `get_config`, capable of instantiating the same
+layer from the config dictionary. It does not handle layer connectivity (handled
+by Network), nor weights (handled by `set_weights`).
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`config`
+</td>
+<td>
+A Python dictionary, typically the
+output of get_config.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+A layer instance.
+</td>
+</tr>
+
+</table>
+
+<h3 id="get_config"><code>get_config</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L288-L293">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_config()
+</code></pre>
+
+Returns the config of the layer.
+
+A layer config is a Python dictionary (serializable) containing the
+configuration of a layer. The same layer can be reinstantiated later (without
+its trained weights) from this configuration.
+
+The config of a layer does not include connectivity information, nor the layer
+class name. These are handled by `Network` (one layer of abstraction above).
+
+Note that `get_config()` does not guarantee to return a fresh copy of dict every
+time it is called. The callers should make a copy of the returned dict if they
+want to modify it.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+Python dictionary.
+</td>
+</tr>
+
+</table>
+
+<h3 id="get_weights"><code>get_weights</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_weights()
+</code></pre>
+
+Returns the current weights of the layer.
+
+The weights of a layer represent the state of the layer. This function returns
+both trainable and non-trainable weight values associated with this layer as a
+list of Numpy arrays, which can in turn be used to load state into similarly
+parameterized layers.
+
+For example, a Dense layer returns a list of two values-- per-output weights and
+the bias value. These can be used to set the weights of another Dense layer:
+
+```
+>>> a = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(1.))
+>>> a_out = a(tf.convert_to_tensor([[1., 2., 3.]]))
+>>> a.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+>>> b = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(2.))
+>>> b_out = b(tf.convert_to_tensor([[10., 20., 30.]]))
+>>> b.get_weights()
+[array([[2.],
+       [2.],
+       [2.]], dtype=float32), array([0.], dtype=float32)]
+>>> b.set_weights(a.get_weights())
+>>> b.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+Weights values as a list of numpy arrays.
+</td>
+</tr>
+
+</table>
+
+<h3 id="set_weights"><code>set_weights</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>set_weights(
+    weights
+)
+</code></pre>
+
+Sets the weights of the layer, from Numpy arrays.
+
+The weights of a layer represent the state of the layer. This function sets the
+weight values from numpy arrays. The weight values should be passed in the order
+they are created by the layer. Note that the layer's weights must be
+instantiated before calling this function by calling the layer.
+
+For example, a Dense layer returns a list of two values-- per-output weights and
+the bias value. These can be used to set the weights of another Dense layer:
+
+```
+>>> a = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(1.))
+>>> a_out = a(tf.convert_to_tensor([[1., 2., 3.]]))
+>>> a.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+>>> b = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(2.))
+>>> b_out = b(tf.convert_to_tensor([[10., 20., 30.]]))
+>>> b.get_weights()
+[array([[2.],
+       [2.],
+       [2.]], dtype=float32), array([0.], dtype=float32)]
+>>> b.set_weights(a.get_weights())
+>>> b.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`weights`
+</td>
+<td>
+a list of Numpy arrays. The number
+of arrays and their shape must match
+number of the dimensions of the weights
+of the layer (i.e. it should match the
+output of `get_weights`).
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Raises</th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+If the provided weights list does not match the
+layer's specifications.
+</td>
+</tr>
+</table>
+
+<h3 id="with_name_scope"><code>with_name_scope</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@classmethod</code>
+<code>with_name_scope(
+    method
+)
+</code></pre>
+
+Decorator to automatically enter the module name scope.
+
+```
+>>> class MyModule(tf.Module):
+...   @tf.Module.with_name_scope
+...   def __call__(self, x):
+...     if not hasattr(self, 'w'):
+...       self.w = tf.Variable(tf.random.normal([x.shape[1], 3]))
+...     return tf.matmul(x, self.w)
+```
+
+Using the above module would produce `tf.Variable`s and `tf.Tensor`s whose names
+included the module name:
+
+```
+>>> mod = MyModule()
+>>> mod(tf.ones([1, 2]))
+<tf.Tensor: shape=(1, 3), dtype=float32, numpy=..., dtype=float32)>
+>>> mod.w
+<tf.Variable 'my_module/Variable:0' shape=(2, 3) dtype=float32,
+numpy=..., dtype=float32)>
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`method`
+</td>
+<td>
+The method to wrap.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+The original method wrapped such that it enters the module's name scope.
+</td>
+</tr>
+
+</table>
+
+<h3 id="__call__"><code>__call__</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>__call__(
+    *args, **kwargs
+)
+</code></pre>
+
+Wraps `call`, applying pre- and post-processing steps.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`*args`
+</td>
+<td>
+Positional arguments to be passed to `self.call`.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Keyword arguments to be passed to `self.call`.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+Output tensor(s).
+</td>
+</tr>
+
+</table>
+
+#### Note:
+
+-   The following optional keyword arguments are reserved for specific uses:
+    *   `training`: Boolean scalar tensor of Python boolean indicating whether
+        the `call` is meant for training or inference.
+    *   `mask`: Boolean input mask.
+-   If the layer's `call` method takes a `mask` argument (as some Keras layers
+    do), its default value will be set to the mask generated for `inputs` by the
+    previous layer (if `input` did come from a layer that generated a
+    corresponding mask, i.e. if it came from a Keras layer with masking support.
+-   If the layer is not built, the method will call `build`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Raises</th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+if the layer's `call` method returns None (an invalid value).
+</td>
+</tr><tr>
+<td>
+`RuntimeError`
+</td>
+<td>
+if `super().__init__()` was not called in the constructor.
+</td>
+</tr>
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/DocumentInteractionAttention.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/DocumentInteractionAttention.md
new file mode 100644
index 0000000..5ef72c8
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/DocumentInteractionAttention.md
@@ -0,0 +1,985 @@
+description: Cross Document Interaction Attention layer.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.layers.DocumentInteractionAttention" />
+<meta itemprop="path" content="Stable" />
+<meta itemprop="property" content="__call__"/>
+<meta itemprop="property" content="__init__"/>
+<meta itemprop="property" content="__new__"/>
+<meta itemprop="property" content="add_loss"/>
+<meta itemprop="property" content="add_metric"/>
+<meta itemprop="property" content="build"/>
+<meta itemprop="property" content="compute_mask"/>
+<meta itemprop="property" content="compute_output_shape"/>
+<meta itemprop="property" content="count_params"/>
+<meta itemprop="property" content="from_config"/>
+<meta itemprop="property" content="get_config"/>
+<meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="set_weights"/>
+<meta itemprop="property" content="with_name_scope"/>
+</div>
+
+# tfr.keras.layers.DocumentInteractionAttention
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L297-L426">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Cross Document Interaction Attention layer.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.layers.DocumentInteractionAttention(
+    num_heads: int,
+    head_size: int,
+    num_layers: int = 1,
+    dropout_rate: float = 0.5,
+    name: Optional[str] = None,
+    **kwargs
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+
+This layer implements the cross-document attention described in "Permutation
+Equivariant Document Interaction Network for Neural Learning to Rank".
+http://research.google/pubs/pub49364/
+
+This layer comprises of several layers of Multi-Headed Attention (MHA) applied
+over the list of documents to attend over itself, using a mask to specify valid
+documents. Optionally, the user can specify the `topk` documents as the listwise
+context that is used to compute the attention per document. If `topk` is None,
+all the documents are used as listwise context to compute the attention.
+
+MHA uses scaled dot product attention, with residual connection and layer
+normalization as follows. This transformation is applied for `num_layers` times:
+h_i := LayerNorm_i(h_{i-1} + MHA_i(h_{i-1}), TopK(h_{i-1}; k))
+
+#### Example:
+
+```python
+  # Batch size = 2, list_size = 3.
+  inputs =  [[[1., 1.], [1., 0.], [1., 1.]], [[0., 0.], [0., 0.], [0., 0.]]]
+  list_mask = [[True, True, False], [True, False, False]]
+  dia_layer = DocumentInteractionAttention(
+      num_heads=1, head_size=64, num_layers=1, topk=1)
+  dia_output = dia_layer(
+      inputs=inputs,
+      training=False,
+      list_mask=list_mask)
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`num_heads`
+</td>
+<td>
+Number of attention heads (see `MultiHeadAttention` for more
+details on this argument).
+</td>
+</tr><tr>
+<td>
+`head_size`
+</td>
+<td>
+Size of each attention head.
+</td>
+</tr><tr>
+<td>
+`num_layers`
+</td>
+<td>
+Number of cross-document attention layers.
+</td>
+</tr><tr>
+<td>
+`dropout_rate`
+</td>
+<td>
+Dropout probability.
+</td>
+</tr><tr>
+<td>
+`name`
+</td>
+<td>
+Name of the layer.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+keyword arguments.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
+
+<tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `input` </td> <td> Retrieves the input
+tensor(s) of a layer.
+
+Only applicable if the layer has exactly one input, i.e. if it is connected to
+one incoming layer. </td> </tr><tr> <td> `input_spec` </td> <td> `InputSpec`
+instance(s) describing the input format for this layer.
+
+When you create a layer subclass, you can set `self.input_spec` to enable the
+layer to run input compatibility checks when it is called. Consider a `Conv2D`
+layer: it can only be called on a single input tensor of rank 4. As such, you
+can set, in `__init__()`:
+
+```python
+self.input_spec = tf.keras.layers.InputSpec(ndim=4)
+```
+
+Now, if you try to call the layer on an input that isn't rank 4 (for instance,
+an input of shape `(2,)`, it will raise a nicely-formatted error:
+
+```
+ValueError: Input 0 of layer conv2d is incompatible with the layer:
+expected ndim=4, found ndim=1. Full shape received: [2]
+```
+
+Input checks that can be specified via `input_spec` include: - Structure (e.g. a
+single input, a list of 2 inputs, etc) - Shape - Rank (ndim) - Dtype
+
+For more information, see `tf.keras.layers.InputSpec`. </td> </tr><tr> <td>
+`losses` </td> <td> List of losses added using the `add_loss()` API.
+
+Variable regularization tensors are created when this property is accessed, so
+it is eager safe: accessing `losses` under a `tf.GradientTape` will propagate
+gradients back to the corresponding variables.
+
+```
+>>> class MyLayer(tf.keras.layers.Layer):
+...   def call(self, inputs):
+...     self.add_loss(tf.abs(tf.reduce_mean(inputs)))
+...     return inputs
+>>> l = MyLayer()
+>>> l(np.ones((10, 1)))
+>>> l.losses
+[1.0]
+```
+
+```
+>>> inputs = tf.keras.Input(shape=(10,))
+>>> x = tf.keras.layers.Dense(10)(inputs)
+>>> outputs = tf.keras.layers.Dense(1)(x)
+>>> model = tf.keras.Model(inputs, outputs)
+>>> # Activity regularization.
+>>> len(model.losses)
+0
+>>> model.add_loss(tf.abs(tf.reduce_mean(x)))
+>>> len(model.losses)
+1
+```
+
+```
+>>> inputs = tf.keras.Input(shape=(10,))
+>>> d = tf.keras.layers.Dense(10, kernel_initializer='ones')
+>>> x = d(inputs)
+>>> outputs = tf.keras.layers.Dense(1)(x)
+>>> model = tf.keras.Model(inputs, outputs)
+>>> # Weight regularization.
+>>> model.add_loss(lambda: tf.reduce_mean(d.kernel))
+>>> model.losses
+[<tf.Tensor: shape=(), dtype=float32, numpy=1.0>]
+```
+
+</td> </tr><tr> <td> `metrics` </td> <td> List of metrics added using the
+`add_metric()` API.
+
+```
+>>> input = tf.keras.layers.Input(shape=(3,))
+>>> d = tf.keras.layers.Dense(2)
+>>> output = d(input)
+>>> d.add_metric(tf.reduce_max(output), name='max')
+>>> d.add_metric(tf.reduce_min(output), name='min')
+>>> [m.name for m in d.metrics]
+['max', 'min']
+```
+
+</td> </tr><tr> <td> `name` </td> <td> Name of the layer (string), set in the
+constructor. </td> </tr><tr> <td> `name_scope` </td> <td> Returns a
+`tf.name_scope` instance for this class. </td> </tr><tr> <td>
+`non_trainable_weights` </td> <td> List of all non-trainable weights tracked by
+this layer.
+
+Non-trainable weights are *not* updated during training. They are expected to be
+updated manually in `call()`. </td> </tr><tr> <td> `output` </td> <td> Retrieves
+the output tensor(s) of a layer.
+
+Only applicable if the layer has exactly one output, i.e. if it is connected to
+one incoming layer. </td> </tr><tr> <td> `submodules` </td> <td> Sequence of all
+sub-modules.
+
+Submodules are modules which are properties of this module, or found as
+properties of modules which are properties of this module (and so on).
+
+```
+>>> a = tf.Module()
+>>> b = tf.Module()
+>>> c = tf.Module()
+>>> a.b = b
+>>> b.c = c
+>>> list(a.submodules) == [b, c]
+True
+>>> list(b.submodules) == [c]
+True
+>>> list(c.submodules) == []
+True
+```
+
+</td> </tr><tr> <td> `supports_masking` </td> <td> Whether this layer supports
+computing a mask using `compute_mask`. </td> </tr><tr> <td> `trainable` </td>
+<td>
+
+</td> </tr><tr> <td> `trainable_weights` </td> <td> List of all trainable
+weights tracked by this layer.
+
+Trainable weights are updated via gradient descent during training.
+</td>
+</tr><tr>
+<td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
+`weights`
+</td>
+<td>
+Returns the list of all layer variables/weights.
+</td>
+</tr>
+</table>
+
+## Methods
+
+<h3 id="add_loss"><code>add_loss</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>add_loss(
+    losses, **kwargs
+)
+</code></pre>
+
+Add loss tensor(s), potentially dependent on layer inputs.
+
+Some losses (for instance, activity regularization losses) may be dependent on
+the inputs passed when calling a layer. Hence, when reusing the same layer on
+different inputs `a` and `b`, some entries in `layer.losses` may be dependent on
+`a` and some on `b`. This method automatically keeps track of dependencies.
+
+This method can be used inside a subclassed layer or model's `call` function, in
+which case `losses` should be a Tensor or list of Tensors.
+
+#### Example:
+
+```python
+class MyLayer(tf.keras.layers.Layer):
+  def call(self, inputs):
+    self.add_loss(tf.abs(tf.reduce_mean(inputs)))
+    return inputs
+```
+
+This method can also be called directly on a Functional Model during
+construction. In this case, any loss Tensors passed to this Model must be
+symbolic and be able to be traced back to the model's `Input`s. These losses
+become part of the model's topology and are tracked in `get_config`.
+
+#### Example:
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+# Activity regularization.
+model.add_loss(tf.abs(tf.reduce_mean(x)))
+```
+
+If this is not the case for your loss (if, for example, your loss references a
+`Variable` of one of the model's layers), you can wrap your loss in a
+zero-argument lambda. These losses are not tracked as part of the model's
+topology since they can't be serialized.
+
+#### Example:
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+d = tf.keras.layers.Dense(10)
+x = d(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+# Weight regularization.
+model.add_loss(lambda: tf.reduce_mean(d.kernel))
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`losses`
+</td>
+<td>
+Loss tensor, or list/tuple of tensors. Rather than tensors, losses
+may also be zero-argument callables which create a loss tensor.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Additional keyword arguments for backward compatibility.
+Accepted values:
+inputs - Deprecated, will be automatically inferred.
+</td>
+</tr>
+</table>
+
+<h3 id="add_metric"><code>add_metric</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>add_metric(
+    value, name=None, **kwargs
+)
+</code></pre>
+
+Adds metric tensor to the layer.
+
+This method can be used inside the `call()` method of a subclassed layer or
+model.
+
+```python
+class MyMetricLayer(tf.keras.layers.Layer):
+  def __init__(self):
+    super(MyMetricLayer, self).__init__(name='my_metric_layer')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
+
+  def call(self, inputs):
+    self.add_metric(self.mean(x))
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
+    return inputs
+```
+
+This method can also be called directly on a Functional Model during
+construction. In this case, any tensor passed to this Model must be symbolic and
+be able to be traced back to the model's `Input`s. These metrics become part of
+the model's topology and are tracked when you save the model via `save()`.
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+model.add_metric(math_ops.reduce_sum(x), name='metric_1')
+```
+
+Note: Calling `add_metric()` with the result of a metric object on a Functional
+Model, as shown in the example below, is not supported. This is because we
+cannot trace the metric result tensor back to the model's inputs.
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+model.add_metric(tf.keras.metrics.Mean()(x), name='metric_1')
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`value`
+</td>
+<td>
+Metric tensor.
+</td>
+</tr><tr>
+<td>
+`name`
+</td>
+<td>
+String metric name.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Additional keyword arguments for backward compatibility.
+Accepted values:
+`aggregation` - When the `value` tensor provided is not the result of
+calling a `keras.Metric` instance, it will be aggregated by default
+using a `keras.Metric.Mean`.
+</td>
+</tr>
+</table>
+
+<h3 id="build"><code>build</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>build(
+    input_shape
+)
+</code></pre>
+
+Creates the variables of the layer (optional, for subclass implementers).
+
+This is a method that implementers of subclasses of `Layer` or `Model` can
+override if they need a state-creation step in-between layer instantiation and
+layer call.
+
+This is typically used to create the weights of `Layer` subclasses.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`input_shape`
+</td>
+<td>
+Instance of `TensorShape`, or list of instances of
+`TensorShape` if the layer expects a list of inputs
+(one instance per input).
+</td>
+</tr>
+</table>
+
+<h3 id="compute_mask"><code>compute_mask</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>compute_mask(
+    inputs, mask=None
+)
+</code></pre>
+
+Computes an output mask tensor.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`inputs`
+</td>
+<td>
+Tensor or list of tensors.
+</td>
+</tr><tr>
+<td>
+`mask`
+</td>
+<td>
+Tensor or list of tensors.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+None or a tensor (or list of tensors,
+one per output tensor of the layer).
+</td>
+</tr>
+
+</table>
+
+<h3 id="compute_output_shape"><code>compute_output_shape</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>compute_output_shape(
+    input_shape
+)
+</code></pre>
+
+Computes the output shape of the layer.
+
+If the layer has not been built, this method will call `build` on the layer.
+This assumes that the layer will later be used with inputs that match the input
+shape provided here.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`input_shape`
+</td>
+<td>
+Shape tuple (tuple of integers)
+or list of shape tuples (one per output tensor of the layer).
+Shape tuples can include None for free dimensions,
+instead of an integer.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+An input shape tuple.
+</td>
+</tr>
+
+</table>
+
+<h3 id="count_params"><code>count_params</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>count_params()
+</code></pre>
+
+Count the total number of scalars composing the weights.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+An integer count.
+</td>
+</tr>
+
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Raises</th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+if the layer isn't yet built
+(in which case its weights aren't yet defined).
+</td>
+</tr>
+</table>
+
+<h3 id="from_config"><code>from_config</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@classmethod</code>
+<code>from_config(
+    config
+)
+</code></pre>
+
+Creates a layer from its config.
+
+This method is the reverse of `get_config`, capable of instantiating the same
+layer from the config dictionary. It does not handle layer connectivity (handled
+by Network), nor weights (handled by `set_weights`).
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`config`
+</td>
+<td>
+A Python dictionary, typically the
+output of get_config.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+A layer instance.
+</td>
+</tr>
+
+</table>
+
+<h3 id="get_config"><code>get_config</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L418-L426">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_config()
+</code></pre>
+
+Returns the config of the layer.
+
+A layer config is a Python dictionary (serializable) containing the
+configuration of a layer. The same layer can be reinstantiated later (without
+its trained weights) from this configuration.
+
+The config of a layer does not include connectivity information, nor the layer
+class name. These are handled by `Network` (one layer of abstraction above).
+
+Note that `get_config()` does not guarantee to return a fresh copy of dict every
+time it is called. The callers should make a copy of the returned dict if they
+want to modify it.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+Python dictionary.
+</td>
+</tr>
+
+</table>
+
+<h3 id="get_weights"><code>get_weights</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_weights()
+</code></pre>
+
+Returns the current weights of the layer.
+
+The weights of a layer represent the state of the layer. This function returns
+both trainable and non-trainable weight values associated with this layer as a
+list of Numpy arrays, which can in turn be used to load state into similarly
+parameterized layers.
+
+For example, a Dense layer returns a list of two values-- per-output weights and
+the bias value. These can be used to set the weights of another Dense layer:
+
+```
+>>> a = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(1.))
+>>> a_out = a(tf.convert_to_tensor([[1., 2., 3.]]))
+>>> a.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+>>> b = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(2.))
+>>> b_out = b(tf.convert_to_tensor([[10., 20., 30.]]))
+>>> b.get_weights()
+[array([[2.],
+       [2.],
+       [2.]], dtype=float32), array([0.], dtype=float32)]
+>>> b.set_weights(a.get_weights())
+>>> b.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+Weights values as a list of numpy arrays.
+</td>
+</tr>
+
+</table>
+
+<h3 id="set_weights"><code>set_weights</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>set_weights(
+    weights
+)
+</code></pre>
+
+Sets the weights of the layer, from Numpy arrays.
+
+The weights of a layer represent the state of the layer. This function sets the
+weight values from numpy arrays. The weight values should be passed in the order
+they are created by the layer. Note that the layer's weights must be
+instantiated before calling this function by calling the layer.
+
+For example, a Dense layer returns a list of two values-- per-output weights and
+the bias value. These can be used to set the weights of another Dense layer:
+
+```
+>>> a = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(1.))
+>>> a_out = a(tf.convert_to_tensor([[1., 2., 3.]]))
+>>> a.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+>>> b = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(2.))
+>>> b_out = b(tf.convert_to_tensor([[10., 20., 30.]]))
+>>> b.get_weights()
+[array([[2.],
+       [2.],
+       [2.]], dtype=float32), array([0.], dtype=float32)]
+>>> b.set_weights(a.get_weights())
+>>> b.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`weights`
+</td>
+<td>
+a list of Numpy arrays. The number
+of arrays and their shape must match
+number of the dimensions of the weights
+of the layer (i.e. it should match the
+output of `get_weights`).
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Raises</th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+If the provided weights list does not match the
+layer's specifications.
+</td>
+</tr>
+</table>
+
+<h3 id="with_name_scope"><code>with_name_scope</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@classmethod</code>
+<code>with_name_scope(
+    method
+)
+</code></pre>
+
+Decorator to automatically enter the module name scope.
+
+```
+>>> class MyModule(tf.Module):
+...   @tf.Module.with_name_scope
+...   def __call__(self, x):
+...     if not hasattr(self, 'w'):
+...       self.w = tf.Variable(tf.random.normal([x.shape[1], 3]))
+...     return tf.matmul(x, self.w)
+```
+
+Using the above module would produce `tf.Variable`s and `tf.Tensor`s whose names
+included the module name:
+
+```
+>>> mod = MyModule()
+>>> mod(tf.ones([1, 2]))
+<tf.Tensor: shape=(1, 3), dtype=float32, numpy=..., dtype=float32)>
+>>> mod.w
+<tf.Variable 'my_module/Variable:0' shape=(2, 3) dtype=float32,
+numpy=..., dtype=float32)>
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`method`
+</td>
+<td>
+The method to wrap.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+The original method wrapped such that it enters the module's name scope.
+</td>
+</tr>
+
+</table>
+
+<h3 id="__call__"><code>__call__</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>__call__(
+    *args, **kwargs
+)
+</code></pre>
+
+Wraps `call`, applying pre- and post-processing steps.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`*args`
+</td>
+<td>
+Positional arguments to be passed to `self.call`.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Keyword arguments to be passed to `self.call`.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+Output tensor(s).
+</td>
+</tr>
+
+</table>
+
+#### Note:
+
+-   The following optional keyword arguments are reserved for specific uses:
+    *   `training`: Boolean scalar tensor of Python boolean indicating whether
+        the `call` is meant for training or inference.
+    *   `mask`: Boolean input mask.
+-   If the layer's `call` method takes a `mask` argument (as some Keras layers
+    do), its default value will be set to the mask generated for `inputs` by the
+    previous layer (if `input` did come from a layer that generated a
+    corresponding mask, i.e. if it came from a Keras layer with masking support.
+-   If the layer is not built, the method will call `build`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Raises</th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+if the layer's `call` method returns None (an invalid value).
+</td>
+</tr><tr>
+<td>
+`RuntimeError`
+</td>
+<td>
+if `super().__init__()` was not called in the constructor.
+</td>
+</tr>
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/FlattenList.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/FlattenList.md
new file mode 100644
index 0000000..f39d2f0
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/FlattenList.md
@@ -0,0 +1,948 @@
+description: Layer to flatten the example list.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.layers.FlattenList" />
+<meta itemprop="path" content="Stable" />
+<meta itemprop="property" content="__call__"/>
+<meta itemprop="property" content="__init__"/>
+<meta itemprop="property" content="__new__"/>
+<meta itemprop="property" content="add_loss"/>
+<meta itemprop="property" content="add_metric"/>
+<meta itemprop="property" content="build"/>
+<meta itemprop="property" content="compute_mask"/>
+<meta itemprop="property" content="compute_output_shape"/>
+<meta itemprop="property" content="count_params"/>
+<meta itemprop="property" content="from_config"/>
+<meta itemprop="property" content="get_config"/>
+<meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="set_weights"/>
+<meta itemprop="property" content="with_name_scope"/>
+</div>
+
+# tfr.keras.layers.FlattenList
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L14-L111">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Layer to flatten the example list.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.layers.FlattenList(
+    circular_padding: bool = True,
+    name: Optional[str] = None,
+    **kwargs
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+
+This layer flattens the batch_size dimension and the list_size dimension for the
+`example_features` and expands list_size times for the `context_features`.
+
+Example use, `python context_features = {'context_feature_1': [[1], [0]]}
+example_features = {'example_feature_1': [[[1], [0], [-1]], [[0], [1], [0]]]}
+mask = [[True, True, False], [True, False, False]] flattened_context_features,
+flattened_example_features = FlattenList()( context_features, example_features,
+mask)` That is, there are two valid examples in the first query and one valid
+example in the second query. Then `python flattened_context_features =
+{'context_feature_1': [[1], [1], [1], [0], [0], [0]]}
+flattened_example_features = {'example_feature_1': [[1], [0], [1], [0], [0],
+[0]]}` `context_feature_1` is repeated by list_size=3 times. `example_feature_1`
+is flattened and padded with the invalid terms replaced by valid terms in each
+query in a circular way.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`circular_padding`
+</td>
+<td>
+Whether to apply circular padding to replace invalid
+features with valid ones.
+</td>
+</tr><tr>
+<td>
+`name`
+</td>
+<td>
+Name of the layer.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+keyword arguments.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
+
+<tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `input` </td> <td> Retrieves the input
+tensor(s) of a layer.
+
+Only applicable if the layer has exactly one input, i.e. if it is connected to
+one incoming layer. </td> </tr><tr> <td> `input_spec` </td> <td> `InputSpec`
+instance(s) describing the input format for this layer.
+
+When you create a layer subclass, you can set `self.input_spec` to enable the
+layer to run input compatibility checks when it is called. Consider a `Conv2D`
+layer: it can only be called on a single input tensor of rank 4. As such, you
+can set, in `__init__()`:
+
+```python
+self.input_spec = tf.keras.layers.InputSpec(ndim=4)
+```
+
+Now, if you try to call the layer on an input that isn't rank 4 (for instance,
+an input of shape `(2,)`, it will raise a nicely-formatted error:
+
+```
+ValueError: Input 0 of layer conv2d is incompatible with the layer:
+expected ndim=4, found ndim=1. Full shape received: [2]
+```
+
+Input checks that can be specified via `input_spec` include: - Structure (e.g. a
+single input, a list of 2 inputs, etc) - Shape - Rank (ndim) - Dtype
+
+For more information, see `tf.keras.layers.InputSpec`. </td> </tr><tr> <td>
+`losses` </td> <td> List of losses added using the `add_loss()` API.
+
+Variable regularization tensors are created when this property is accessed, so
+it is eager safe: accessing `losses` under a `tf.GradientTape` will propagate
+gradients back to the corresponding variables.
+
+```
+>>> class MyLayer(tf.keras.layers.Layer):
+...   def call(self, inputs):
+...     self.add_loss(tf.abs(tf.reduce_mean(inputs)))
+...     return inputs
+>>> l = MyLayer()
+>>> l(np.ones((10, 1)))
+>>> l.losses
+[1.0]
+```
+
+```
+>>> inputs = tf.keras.Input(shape=(10,))
+>>> x = tf.keras.layers.Dense(10)(inputs)
+>>> outputs = tf.keras.layers.Dense(1)(x)
+>>> model = tf.keras.Model(inputs, outputs)
+>>> # Activity regularization.
+>>> len(model.losses)
+0
+>>> model.add_loss(tf.abs(tf.reduce_mean(x)))
+>>> len(model.losses)
+1
+```
+
+```
+>>> inputs = tf.keras.Input(shape=(10,))
+>>> d = tf.keras.layers.Dense(10, kernel_initializer='ones')
+>>> x = d(inputs)
+>>> outputs = tf.keras.layers.Dense(1)(x)
+>>> model = tf.keras.Model(inputs, outputs)
+>>> # Weight regularization.
+>>> model.add_loss(lambda: tf.reduce_mean(d.kernel))
+>>> model.losses
+[<tf.Tensor: shape=(), dtype=float32, numpy=1.0>]
+```
+
+</td> </tr><tr> <td> `metrics` </td> <td> List of metrics added using the
+`add_metric()` API.
+
+```
+>>> input = tf.keras.layers.Input(shape=(3,))
+>>> d = tf.keras.layers.Dense(2)
+>>> output = d(input)
+>>> d.add_metric(tf.reduce_max(output), name='max')
+>>> d.add_metric(tf.reduce_min(output), name='min')
+>>> [m.name for m in d.metrics]
+['max', 'min']
+```
+
+</td> </tr><tr> <td> `name` </td> <td> Name of the layer (string), set in the
+constructor. </td> </tr><tr> <td> `name_scope` </td> <td> Returns a
+`tf.name_scope` instance for this class. </td> </tr><tr> <td>
+`non_trainable_weights` </td> <td> List of all non-trainable weights tracked by
+this layer.
+
+Non-trainable weights are *not* updated during training. They are expected to be
+updated manually in `call()`. </td> </tr><tr> <td> `output` </td> <td> Retrieves
+the output tensor(s) of a layer.
+
+Only applicable if the layer has exactly one output, i.e. if it is connected to
+one incoming layer. </td> </tr><tr> <td> `submodules` </td> <td> Sequence of all
+sub-modules.
+
+Submodules are modules which are properties of this module, or found as
+properties of modules which are properties of this module (and so on).
+
+```
+>>> a = tf.Module()
+>>> b = tf.Module()
+>>> c = tf.Module()
+>>> a.b = b
+>>> b.c = c
+>>> list(a.submodules) == [b, c]
+True
+>>> list(b.submodules) == [c]
+True
+>>> list(c.submodules) == []
+True
+```
+
+</td> </tr><tr> <td> `supports_masking` </td> <td> Whether this layer supports
+computing a mask using `compute_mask`. </td> </tr><tr> <td> `trainable` </td>
+<td>
+
+</td> </tr><tr> <td> `trainable_weights` </td> <td> List of all trainable
+weights tracked by this layer.
+
+Trainable weights are updated via gradient descent during training.
+</td>
+</tr><tr>
+<td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
+`weights`
+</td>
+<td>
+Returns the list of all layer variables/weights.
+</td>
+</tr>
+</table>
+
+## Methods
+
+<h3 id="add_loss"><code>add_loss</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>add_loss(
+    losses, **kwargs
+)
+</code></pre>
+
+Add loss tensor(s), potentially dependent on layer inputs.
+
+Some losses (for instance, activity regularization losses) may be dependent on
+the inputs passed when calling a layer. Hence, when reusing the same layer on
+different inputs `a` and `b`, some entries in `layer.losses` may be dependent on
+`a` and some on `b`. This method automatically keeps track of dependencies.
+
+This method can be used inside a subclassed layer or model's `call` function, in
+which case `losses` should be a Tensor or list of Tensors.
+
+#### Example:
+
+```python
+class MyLayer(tf.keras.layers.Layer):
+  def call(self, inputs):
+    self.add_loss(tf.abs(tf.reduce_mean(inputs)))
+    return inputs
+```
+
+This method can also be called directly on a Functional Model during
+construction. In this case, any loss Tensors passed to this Model must be
+symbolic and be able to be traced back to the model's `Input`s. These losses
+become part of the model's topology and are tracked in `get_config`.
+
+#### Example:
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+# Activity regularization.
+model.add_loss(tf.abs(tf.reduce_mean(x)))
+```
+
+If this is not the case for your loss (if, for example, your loss references a
+`Variable` of one of the model's layers), you can wrap your loss in a
+zero-argument lambda. These losses are not tracked as part of the model's
+topology since they can't be serialized.
+
+#### Example:
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+d = tf.keras.layers.Dense(10)
+x = d(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+# Weight regularization.
+model.add_loss(lambda: tf.reduce_mean(d.kernel))
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`losses`
+</td>
+<td>
+Loss tensor, or list/tuple of tensors. Rather than tensors, losses
+may also be zero-argument callables which create a loss tensor.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Additional keyword arguments for backward compatibility.
+Accepted values:
+inputs - Deprecated, will be automatically inferred.
+</td>
+</tr>
+</table>
+
+<h3 id="add_metric"><code>add_metric</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>add_metric(
+    value, name=None, **kwargs
+)
+</code></pre>
+
+Adds metric tensor to the layer.
+
+This method can be used inside the `call()` method of a subclassed layer or
+model.
+
+```python
+class MyMetricLayer(tf.keras.layers.Layer):
+  def __init__(self):
+    super(MyMetricLayer, self).__init__(name='my_metric_layer')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
+
+  def call(self, inputs):
+    self.add_metric(self.mean(x))
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
+    return inputs
+```
+
+This method can also be called directly on a Functional Model during
+construction. In this case, any tensor passed to this Model must be symbolic and
+be able to be traced back to the model's `Input`s. These metrics become part of
+the model's topology and are tracked when you save the model via `save()`.
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+model.add_metric(math_ops.reduce_sum(x), name='metric_1')
+```
+
+Note: Calling `add_metric()` with the result of a metric object on a Functional
+Model, as shown in the example below, is not supported. This is because we
+cannot trace the metric result tensor back to the model's inputs.
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+model.add_metric(tf.keras.metrics.Mean()(x), name='metric_1')
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`value`
+</td>
+<td>
+Metric tensor.
+</td>
+</tr><tr>
+<td>
+`name`
+</td>
+<td>
+String metric name.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Additional keyword arguments for backward compatibility.
+Accepted values:
+`aggregation` - When the `value` tensor provided is not the result of
+calling a `keras.Metric` instance, it will be aggregated by default
+using a `keras.Metric.Mean`.
+</td>
+</tr>
+</table>
+
+<h3 id="build"><code>build</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>build(
+    input_shape
+)
+</code></pre>
+
+Creates the variables of the layer (optional, for subclass implementers).
+
+This is a method that implementers of subclasses of `Layer` or `Model` can
+override if they need a state-creation step in-between layer instantiation and
+layer call.
+
+This is typically used to create the weights of `Layer` subclasses.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`input_shape`
+</td>
+<td>
+Instance of `TensorShape`, or list of instances of
+`TensorShape` if the layer expects a list of inputs
+(one instance per input).
+</td>
+</tr>
+</table>
+
+<h3 id="compute_mask"><code>compute_mask</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>compute_mask(
+    inputs, mask=None
+)
+</code></pre>
+
+Computes an output mask tensor.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`inputs`
+</td>
+<td>
+Tensor or list of tensors.
+</td>
+</tr><tr>
+<td>
+`mask`
+</td>
+<td>
+Tensor or list of tensors.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+None or a tensor (or list of tensors,
+one per output tensor of the layer).
+</td>
+</tr>
+
+</table>
+
+<h3 id="compute_output_shape"><code>compute_output_shape</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>compute_output_shape(
+    input_shape
+)
+</code></pre>
+
+Computes the output shape of the layer.
+
+If the layer has not been built, this method will call `build` on the layer.
+This assumes that the layer will later be used with inputs that match the input
+shape provided here.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`input_shape`
+</td>
+<td>
+Shape tuple (tuple of integers)
+or list of shape tuples (one per output tensor of the layer).
+Shape tuples can include None for free dimensions,
+instead of an integer.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+An input shape tuple.
+</td>
+</tr>
+
+</table>
+
+<h3 id="count_params"><code>count_params</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>count_params()
+</code></pre>
+
+Count the total number of scalars composing the weights.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+An integer count.
+</td>
+</tr>
+
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Raises</th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+if the layer isn't yet built
+(in which case its weights aren't yet defined).
+</td>
+</tr>
+</table>
+
+<h3 id="from_config"><code>from_config</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@classmethod</code>
+<code>from_config(
+    config
+)
+</code></pre>
+
+Creates a layer from its config.
+
+This method is the reverse of `get_config`, capable of instantiating the same
+layer from the config dictionary. It does not handle layer connectivity (handled
+by Network), nor weights (handled by `set_weights`).
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`config`
+</td>
+<td>
+A Python dictionary, typically the
+output of get_config.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+A layer instance.
+</td>
+</tr>
+
+</table>
+
+<h3 id="get_config"><code>get_config</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L106-L111">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_config()
+</code></pre>
+
+Returns the config of the layer.
+
+A layer config is a Python dictionary (serializable) containing the
+configuration of a layer. The same layer can be reinstantiated later (without
+its trained weights) from this configuration.
+
+The config of a layer does not include connectivity information, nor the layer
+class name. These are handled by `Network` (one layer of abstraction above).
+
+Note that `get_config()` does not guarantee to return a fresh copy of dict every
+time it is called. The callers should make a copy of the returned dict if they
+want to modify it.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+Python dictionary.
+</td>
+</tr>
+
+</table>
+
+<h3 id="get_weights"><code>get_weights</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_weights()
+</code></pre>
+
+Returns the current weights of the layer.
+
+The weights of a layer represent the state of the layer. This function returns
+both trainable and non-trainable weight values associated with this layer as a
+list of Numpy arrays, which can in turn be used to load state into similarly
+parameterized layers.
+
+For example, a Dense layer returns a list of two values-- per-output weights and
+the bias value. These can be used to set the weights of another Dense layer:
+
+```
+>>> a = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(1.))
+>>> a_out = a(tf.convert_to_tensor([[1., 2., 3.]]))
+>>> a.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+>>> b = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(2.))
+>>> b_out = b(tf.convert_to_tensor([[10., 20., 30.]]))
+>>> b.get_weights()
+[array([[2.],
+       [2.],
+       [2.]], dtype=float32), array([0.], dtype=float32)]
+>>> b.set_weights(a.get_weights())
+>>> b.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+Weights values as a list of numpy arrays.
+</td>
+</tr>
+
+</table>
+
+<h3 id="set_weights"><code>set_weights</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>set_weights(
+    weights
+)
+</code></pre>
+
+Sets the weights of the layer, from Numpy arrays.
+
+The weights of a layer represent the state of the layer. This function sets the
+weight values from numpy arrays. The weight values should be passed in the order
+they are created by the layer. Note that the layer's weights must be
+instantiated before calling this function by calling the layer.
+
+For example, a Dense layer returns a list of two values-- per-output weights and
+the bias value. These can be used to set the weights of another Dense layer:
+
+```
+>>> a = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(1.))
+>>> a_out = a(tf.convert_to_tensor([[1., 2., 3.]]))
+>>> a.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+>>> b = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(2.))
+>>> b_out = b(tf.convert_to_tensor([[10., 20., 30.]]))
+>>> b.get_weights()
+[array([[2.],
+       [2.],
+       [2.]], dtype=float32), array([0.], dtype=float32)]
+>>> b.set_weights(a.get_weights())
+>>> b.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`weights`
+</td>
+<td>
+a list of Numpy arrays. The number
+of arrays and their shape must match
+number of the dimensions of the weights
+of the layer (i.e. it should match the
+output of `get_weights`).
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Raises</th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+If the provided weights list does not match the
+layer's specifications.
+</td>
+</tr>
+</table>
+
+<h3 id="with_name_scope"><code>with_name_scope</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@classmethod</code>
+<code>with_name_scope(
+    method
+)
+</code></pre>
+
+Decorator to automatically enter the module name scope.
+
+```
+>>> class MyModule(tf.Module):
+...   @tf.Module.with_name_scope
+...   def __call__(self, x):
+...     if not hasattr(self, 'w'):
+...       self.w = tf.Variable(tf.random.normal([x.shape[1], 3]))
+...     return tf.matmul(x, self.w)
+```
+
+Using the above module would produce `tf.Variable`s and `tf.Tensor`s whose names
+included the module name:
+
+```
+>>> mod = MyModule()
+>>> mod(tf.ones([1, 2]))
+<tf.Tensor: shape=(1, 3), dtype=float32, numpy=..., dtype=float32)>
+>>> mod.w
+<tf.Variable 'my_module/Variable:0' shape=(2, 3) dtype=float32,
+numpy=..., dtype=float32)>
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`method`
+</td>
+<td>
+The method to wrap.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+The original method wrapped such that it enters the module's name scope.
+</td>
+</tr>
+
+</table>
+
+<h3 id="__call__"><code>__call__</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>__call__(
+    *args, **kwargs
+)
+</code></pre>
+
+Wraps `call`, applying pre- and post-processing steps.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`*args`
+</td>
+<td>
+Positional arguments to be passed to `self.call`.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Keyword arguments to be passed to `self.call`.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+Output tensor(s).
+</td>
+</tr>
+
+</table>
+
+#### Note:
+
+-   The following optional keyword arguments are reserved for specific uses:
+    *   `training`: Boolean scalar tensor of Python boolean indicating whether
+        the `call` is meant for training or inference.
+    *   `mask`: Boolean input mask.
+-   If the layer's `call` method takes a `mask` argument (as some Keras layers
+    do), its default value will be set to the mask generated for `inputs` by the
+    previous layer (if `input` did come from a layer that generated a
+    corresponding mask, i.e. if it came from a Keras layer with masking support.
+-   If the layer is not built, the method will call `build`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Raises</th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+if the layer's `call` method returns None (an invalid value).
+</td>
+</tr><tr>
+<td>
+`RuntimeError`
+</td>
+<td>
+if `super().__init__()` was not called in the constructor.
+</td>
+</tr>
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/RestoreList.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/RestoreList.md
new file mode 100644
index 0000000..7882eb2
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/layers/RestoreList.md
@@ -0,0 +1,923 @@
+description: Output layer to restore listwise output shape.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.layers.RestoreList" />
+<meta itemprop="path" content="Stable" />
+<meta itemprop="property" content="__call__"/>
+<meta itemprop="property" content="__init__"/>
+<meta itemprop="property" content="__new__"/>
+<meta itemprop="property" content="add_loss"/>
+<meta itemprop="property" content="add_metric"/>
+<meta itemprop="property" content="build"/>
+<meta itemprop="property" content="compute_mask"/>
+<meta itemprop="property" content="compute_output_shape"/>
+<meta itemprop="property" content="count_params"/>
+<meta itemprop="property" content="from_config"/>
+<meta itemprop="property" content="get_config"/>
+<meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="set_weights"/>
+<meta itemprop="property" content="with_name_scope"/>
+</div>
+
+# tfr.keras.layers.RestoreList
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L115-L199">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Output layer to restore listwise output shape.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.layers.RestoreList(
+    name: Optional[str] = None,
+    by_scatter: bool = False,
+    **kwargs
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+
+This layer maps the flattened 1D logits or 2D tensor of shape [batch_size *
+list_size, 1] back to 2D of shape [batch_size, list_size] and mask the invalid
+terms to be a defined large negative value.
+
+Example use, `python flattened_logits = [1, 0.5, 2, 0, -1, 0] mask = [[True,
+True, False], [True, False, False]] logits = RestoreList()(flattened_logits,
+mask)` Then `python logits = [[1, 0.5, log(_EPSILON)], [0, log(_EPSILON),
+log(_EPSILON)]]` where _EPSILON=1e-10. This layer works also for 2D
+`flattened_logits` like [[1], [0.5], [2], [0], [-1], [0]].
+
+When `by_scatter=True`, an nd_indices will be generated using `mask` in the same
+way as `FlattenList`. All values in the `flattened_logits` will be used and
+repeated entries will be averaged. `python flattened_logits = [1, 0.5, 2, 0, -1,
+0] mask = [[True, True, False], [True, False, False]] logits =
+RestoreList(by_scatter=True)(flattened_logits, mask)` Then `python logits =
+[[1.5, 0.5, log(_EPSILON)], [-1/3, log(_EPSILON), log(_EPSILON)]]` This is
+because the flattened_logits are treated as circularly padded entries. The [1st,
+3rd] values [1, 2] are counted to logits[0, 0]. The [4th, 5th, 6th] values [0,
+-1, 0] are counted to logits[1, 0]. Note that We use different values for those
+repeated entries, while they are likely the same in practice.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
+
+<tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `input` </td> <td> Retrieves the input
+tensor(s) of a layer.
+
+Only applicable if the layer has exactly one input, i.e. if it is connected to
+one incoming layer. </td> </tr><tr> <td> `input_spec` </td> <td> `InputSpec`
+instance(s) describing the input format for this layer.
+
+When you create a layer subclass, you can set `self.input_spec` to enable the
+layer to run input compatibility checks when it is called. Consider a `Conv2D`
+layer: it can only be called on a single input tensor of rank 4. As such, you
+can set, in `__init__()`:
+
+```python
+self.input_spec = tf.keras.layers.InputSpec(ndim=4)
+```
+
+Now, if you try to call the layer on an input that isn't rank 4 (for instance,
+an input of shape `(2,)`, it will raise a nicely-formatted error:
+
+```
+ValueError: Input 0 of layer conv2d is incompatible with the layer:
+expected ndim=4, found ndim=1. Full shape received: [2]
+```
+
+Input checks that can be specified via `input_spec` include: - Structure (e.g. a
+single input, a list of 2 inputs, etc) - Shape - Rank (ndim) - Dtype
+
+For more information, see `tf.keras.layers.InputSpec`. </td> </tr><tr> <td>
+`losses` </td> <td> List of losses added using the `add_loss()` API.
+
+Variable regularization tensors are created when this property is accessed, so
+it is eager safe: accessing `losses` under a `tf.GradientTape` will propagate
+gradients back to the corresponding variables.
+
+```
+>>> class MyLayer(tf.keras.layers.Layer):
+...   def call(self, inputs):
+...     self.add_loss(tf.abs(tf.reduce_mean(inputs)))
+...     return inputs
+>>> l = MyLayer()
+>>> l(np.ones((10, 1)))
+>>> l.losses
+[1.0]
+```
+
+```
+>>> inputs = tf.keras.Input(shape=(10,))
+>>> x = tf.keras.layers.Dense(10)(inputs)
+>>> outputs = tf.keras.layers.Dense(1)(x)
+>>> model = tf.keras.Model(inputs, outputs)
+>>> # Activity regularization.
+>>> len(model.losses)
+0
+>>> model.add_loss(tf.abs(tf.reduce_mean(x)))
+>>> len(model.losses)
+1
+```
+
+```
+>>> inputs = tf.keras.Input(shape=(10,))
+>>> d = tf.keras.layers.Dense(10, kernel_initializer='ones')
+>>> x = d(inputs)
+>>> outputs = tf.keras.layers.Dense(1)(x)
+>>> model = tf.keras.Model(inputs, outputs)
+>>> # Weight regularization.
+>>> model.add_loss(lambda: tf.reduce_mean(d.kernel))
+>>> model.losses
+[<tf.Tensor: shape=(), dtype=float32, numpy=1.0>]
+```
+
+</td> </tr><tr> <td> `metrics` </td> <td> List of metrics added using the
+`add_metric()` API.
+
+```
+>>> input = tf.keras.layers.Input(shape=(3,))
+>>> d = tf.keras.layers.Dense(2)
+>>> output = d(input)
+>>> d.add_metric(tf.reduce_max(output), name='max')
+>>> d.add_metric(tf.reduce_min(output), name='min')
+>>> [m.name for m in d.metrics]
+['max', 'min']
+```
+
+</td> </tr><tr> <td> `name` </td> <td> Name of the layer (string), set in the
+constructor. </td> </tr><tr> <td> `name_scope` </td> <td> Returns a
+`tf.name_scope` instance for this class. </td> </tr><tr> <td>
+`non_trainable_weights` </td> <td> List of all non-trainable weights tracked by
+this layer.
+
+Non-trainable weights are *not* updated during training. They are expected to be
+updated manually in `call()`. </td> </tr><tr> <td> `output` </td> <td> Retrieves
+the output tensor(s) of a layer.
+
+Only applicable if the layer has exactly one output, i.e. if it is connected to
+one incoming layer. </td> </tr><tr> <td> `submodules` </td> <td> Sequence of all
+sub-modules.
+
+Submodules are modules which are properties of this module, or found as
+properties of modules which are properties of this module (and so on).
+
+```
+>>> a = tf.Module()
+>>> b = tf.Module()
+>>> c = tf.Module()
+>>> a.b = b
+>>> b.c = c
+>>> list(a.submodules) == [b, c]
+True
+>>> list(b.submodules) == [c]
+True
+>>> list(c.submodules) == []
+True
+```
+
+</td> </tr><tr> <td> `supports_masking` </td> <td> Whether this layer supports
+computing a mask using `compute_mask`. </td> </tr><tr> <td> `trainable` </td>
+<td>
+
+</td> </tr><tr> <td> `trainable_weights` </td> <td> List of all trainable
+weights tracked by this layer.
+
+Trainable weights are updated via gradient descent during training.
+</td>
+</tr><tr>
+<td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
+`weights`
+</td>
+<td>
+Returns the list of all layer variables/weights.
+</td>
+</tr>
+</table>
+
+## Methods
+
+<h3 id="add_loss"><code>add_loss</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>add_loss(
+    losses, **kwargs
+)
+</code></pre>
+
+Add loss tensor(s), potentially dependent on layer inputs.
+
+Some losses (for instance, activity regularization losses) may be dependent on
+the inputs passed when calling a layer. Hence, when reusing the same layer on
+different inputs `a` and `b`, some entries in `layer.losses` may be dependent on
+`a` and some on `b`. This method automatically keeps track of dependencies.
+
+This method can be used inside a subclassed layer or model's `call` function, in
+which case `losses` should be a Tensor or list of Tensors.
+
+#### Example:
+
+```python
+class MyLayer(tf.keras.layers.Layer):
+  def call(self, inputs):
+    self.add_loss(tf.abs(tf.reduce_mean(inputs)))
+    return inputs
+```
+
+This method can also be called directly on a Functional Model during
+construction. In this case, any loss Tensors passed to this Model must be
+symbolic and be able to be traced back to the model's `Input`s. These losses
+become part of the model's topology and are tracked in `get_config`.
+
+#### Example:
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+# Activity regularization.
+model.add_loss(tf.abs(tf.reduce_mean(x)))
+```
+
+If this is not the case for your loss (if, for example, your loss references a
+`Variable` of one of the model's layers), you can wrap your loss in a
+zero-argument lambda. These losses are not tracked as part of the model's
+topology since they can't be serialized.
+
+#### Example:
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+d = tf.keras.layers.Dense(10)
+x = d(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+# Weight regularization.
+model.add_loss(lambda: tf.reduce_mean(d.kernel))
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`losses`
+</td>
+<td>
+Loss tensor, or list/tuple of tensors. Rather than tensors, losses
+may also be zero-argument callables which create a loss tensor.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Additional keyword arguments for backward compatibility.
+Accepted values:
+inputs - Deprecated, will be automatically inferred.
+</td>
+</tr>
+</table>
+
+<h3 id="add_metric"><code>add_metric</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>add_metric(
+    value, name=None, **kwargs
+)
+</code></pre>
+
+Adds metric tensor to the layer.
+
+This method can be used inside the `call()` method of a subclassed layer or
+model.
+
+```python
+class MyMetricLayer(tf.keras.layers.Layer):
+  def __init__(self):
+    super(MyMetricLayer, self).__init__(name='my_metric_layer')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
+
+  def call(self, inputs):
+    self.add_metric(self.mean(x))
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
+    return inputs
+```
+
+This method can also be called directly on a Functional Model during
+construction. In this case, any tensor passed to this Model must be symbolic and
+be able to be traced back to the model's `Input`s. These metrics become part of
+the model's topology and are tracked when you save the model via `save()`.
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+model.add_metric(math_ops.reduce_sum(x), name='metric_1')
+```
+
+Note: Calling `add_metric()` with the result of a metric object on a Functional
+Model, as shown in the example below, is not supported. This is because we
+cannot trace the metric result tensor back to the model's inputs.
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+model.add_metric(tf.keras.metrics.Mean()(x), name='metric_1')
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`value`
+</td>
+<td>
+Metric tensor.
+</td>
+</tr><tr>
+<td>
+`name`
+</td>
+<td>
+String metric name.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Additional keyword arguments for backward compatibility.
+Accepted values:
+`aggregation` - When the `value` tensor provided is not the result of
+calling a `keras.Metric` instance, it will be aggregated by default
+using a `keras.Metric.Mean`.
+</td>
+</tr>
+</table>
+
+<h3 id="build"><code>build</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>build(
+    input_shape
+)
+</code></pre>
+
+Creates the variables of the layer (optional, for subclass implementers).
+
+This is a method that implementers of subclasses of `Layer` or `Model` can
+override if they need a state-creation step in-between layer instantiation and
+layer call.
+
+This is typically used to create the weights of `Layer` subclasses.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`input_shape`
+</td>
+<td>
+Instance of `TensorShape`, or list of instances of
+`TensorShape` if the layer expects a list of inputs
+(one instance per input).
+</td>
+</tr>
+</table>
+
+<h3 id="compute_mask"><code>compute_mask</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>compute_mask(
+    inputs, mask=None
+)
+</code></pre>
+
+Computes an output mask tensor.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`inputs`
+</td>
+<td>
+Tensor or list of tensors.
+</td>
+</tr><tr>
+<td>
+`mask`
+</td>
+<td>
+Tensor or list of tensors.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+None or a tensor (or list of tensors,
+one per output tensor of the layer).
+</td>
+</tr>
+
+</table>
+
+<h3 id="compute_output_shape"><code>compute_output_shape</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>compute_output_shape(
+    input_shape
+)
+</code></pre>
+
+Computes the output shape of the layer.
+
+If the layer has not been built, this method will call `build` on the layer.
+This assumes that the layer will later be used with inputs that match the input
+shape provided here.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`input_shape`
+</td>
+<td>
+Shape tuple (tuple of integers)
+or list of shape tuples (one per output tensor of the layer).
+Shape tuples can include None for free dimensions,
+instead of an integer.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+An input shape tuple.
+</td>
+</tr>
+
+</table>
+
+<h3 id="count_params"><code>count_params</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>count_params()
+</code></pre>
+
+Count the total number of scalars composing the weights.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+An integer count.
+</td>
+</tr>
+
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Raises</th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+if the layer isn't yet built
+(in which case its weights aren't yet defined).
+</td>
+</tr>
+</table>
+
+<h3 id="from_config"><code>from_config</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@classmethod</code>
+<code>from_config(
+    config
+)
+</code></pre>
+
+Creates a layer from its config.
+
+This method is the reverse of `get_config`, capable of instantiating the same
+layer from the config dictionary. It does not handle layer connectivity (handled
+by Network), nor weights (handled by `set_weights`).
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`config`
+</td>
+<td>
+A Python dictionary, typically the
+output of get_config.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+A layer instance.
+</td>
+</tr>
+
+</table>
+
+<h3 id="get_config"><code>get_config</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L194-L199">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_config()
+</code></pre>
+
+Returns the config of the layer.
+
+A layer config is a Python dictionary (serializable) containing the
+configuration of a layer. The same layer can be reinstantiated later (without
+its trained weights) from this configuration.
+
+The config of a layer does not include connectivity information, nor the layer
+class name. These are handled by `Network` (one layer of abstraction above).
+
+Note that `get_config()` does not guarantee to return a fresh copy of dict every
+time it is called. The callers should make a copy of the returned dict if they
+want to modify it.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+Python dictionary.
+</td>
+</tr>
+
+</table>
+
+<h3 id="get_weights"><code>get_weights</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_weights()
+</code></pre>
+
+Returns the current weights of the layer.
+
+The weights of a layer represent the state of the layer. This function returns
+both trainable and non-trainable weight values associated with this layer as a
+list of Numpy arrays, which can in turn be used to load state into similarly
+parameterized layers.
+
+For example, a Dense layer returns a list of two values-- per-output weights and
+the bias value. These can be used to set the weights of another Dense layer:
+
+```
+>>> a = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(1.))
+>>> a_out = a(tf.convert_to_tensor([[1., 2., 3.]]))
+>>> a.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+>>> b = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(2.))
+>>> b_out = b(tf.convert_to_tensor([[10., 20., 30.]]))
+>>> b.get_weights()
+[array([[2.],
+       [2.],
+       [2.]], dtype=float32), array([0.], dtype=float32)]
+>>> b.set_weights(a.get_weights())
+>>> b.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+Weights values as a list of numpy arrays.
+</td>
+</tr>
+
+</table>
+
+<h3 id="set_weights"><code>set_weights</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>set_weights(
+    weights
+)
+</code></pre>
+
+Sets the weights of the layer, from Numpy arrays.
+
+The weights of a layer represent the state of the layer. This function sets the
+weight values from numpy arrays. The weight values should be passed in the order
+they are created by the layer. Note that the layer's weights must be
+instantiated before calling this function by calling the layer.
+
+For example, a Dense layer returns a list of two values-- per-output weights and
+the bias value. These can be used to set the weights of another Dense layer:
+
+```
+>>> a = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(1.))
+>>> a_out = a(tf.convert_to_tensor([[1., 2., 3.]]))
+>>> a.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+>>> b = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(2.))
+>>> b_out = b(tf.convert_to_tensor([[10., 20., 30.]]))
+>>> b.get_weights()
+[array([[2.],
+       [2.],
+       [2.]], dtype=float32), array([0.], dtype=float32)]
+>>> b.set_weights(a.get_weights())
+>>> b.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`weights`
+</td>
+<td>
+a list of Numpy arrays. The number
+of arrays and their shape must match
+number of the dimensions of the weights
+of the layer (i.e. it should match the
+output of `get_weights`).
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Raises</th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+If the provided weights list does not match the
+layer's specifications.
+</td>
+</tr>
+</table>
+
+<h3 id="with_name_scope"><code>with_name_scope</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@classmethod</code>
+<code>with_name_scope(
+    method
+)
+</code></pre>
+
+Decorator to automatically enter the module name scope.
+
+```
+>>> class MyModule(tf.Module):
+...   @tf.Module.with_name_scope
+...   def __call__(self, x):
+...     if not hasattr(self, 'w'):
+...       self.w = tf.Variable(tf.random.normal([x.shape[1], 3]))
+...     return tf.matmul(x, self.w)
+```
+
+Using the above module would produce `tf.Variable`s and `tf.Tensor`s whose names
+included the module name:
+
+```
+>>> mod = MyModule()
+>>> mod(tf.ones([1, 2]))
+<tf.Tensor: shape=(1, 3), dtype=float32, numpy=..., dtype=float32)>
+>>> mod.w
+<tf.Variable 'my_module/Variable:0' shape=(2, 3) dtype=float32,
+numpy=..., dtype=float32)>
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`method`
+</td>
+<td>
+The method to wrap.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+The original method wrapped such that it enters the module's name scope.
+</td>
+</tr>
+
+</table>
+
+<h3 id="__call__"><code>__call__</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>__call__(
+    *args, **kwargs
+)
+</code></pre>
+
+Wraps `call`, applying pre- and post-processing steps.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`*args`
+</td>
+<td>
+Positional arguments to be passed to `self.call`.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Keyword arguments to be passed to `self.call`.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+Output tensor(s).
+</td>
+</tr>
+
+</table>
+
+#### Note:
+
+-   The following optional keyword arguments are reserved for specific uses:
+    *   `training`: Boolean scalar tensor of Python boolean indicating whether
+        the `call` is meant for training or inference.
+    *   `mask`: Boolean input mask.
+-   If the layer's `call` method takes a `mask` argument (as some Keras layers
+    do), its default value will be set to the mask generated for `inputs` by the
+    previous layer (if `input` did come from a layer that generated a
+    corresponding mask, i.e. if it came from a Keras layer with masking support.
+-   If the layer is not built, the method will call `build`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Raises</th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+if the layer's `call` method returns None (an invalid value).
+</td>
+</tr><tr>
+<td>
+`RuntimeError`
+</td>
+<td>
+if `super().__init__()` was not called in the constructor.
+</td>
+</tr>
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses.md
index 923088c..7d724ca 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses.md
@@ -1,3 +1,5 @@
+description: Keras losses in TF-Ranking.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.losses" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
@@ -27,14 +28,25 @@ approximate MRR loss.
 [`class ApproxNDCGLoss`](../../tfr/keras/losses/ApproxNDCGLoss.md): For
 approximate NDCG loss.
 
+[`class ClickEMLoss`](../../tfr/keras/losses/ClickEMLoss.md): For click EM loss.
+
+[`class DCGLambdaWeight`](../../tfr/keras/losses/DCGLambdaWeight.md): Keras
+serializable class for DCG.
+
 [`class GumbelApproxNDCGLoss`](../../tfr/keras/losses/GumbelApproxNDCGLoss.md):
 For Gumbel approximate NDCG loss.
 
+[`class ListMLELambdaWeight`](../../tfr/keras/losses/ListMLELambdaWeight.md):
+LambdaWeight for ListMLE cost function.
+
 [`class ListMLELoss`](../../tfr/keras/losses/ListMLELoss.md): For List MLE loss.
 
 [`class MeanSquaredLoss`](../../tfr/keras/losses/MeanSquaredLoss.md): For mean
 squared error loss.
 
+[`class NDCGLambdaWeight`](../../tfr/keras/losses/NDCGLambdaWeight.md): Keras
+serializable class for NDCG.
+
 [`class PairwiseHingeLoss`](../../tfr/keras/losses/PairwiseHingeLoss.md): For
 pairwise hinge loss.
 
@@ -44,6 +56,9 @@ For pairwise logistic loss.
 [`class PairwiseSoftZeroOneLoss`](../../tfr/keras/losses/PairwiseSoftZeroOneLoss.md):
 For pairwise soft zero-one loss.
 
+[`class PrecisionLambdaWeight`](../../tfr/keras/losses/PrecisionLambdaWeight.md):
+Keras serializable class for Precision.
+
 [`class RankingLossKey`](../../tfr/keras/losses/RankingLossKey.md): Ranking loss
 key strings.
 
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ApproxMRRLoss.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ApproxMRRLoss.md
index 9b26527..f8c9e3d 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ApproxMRRLoss.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ApproxMRRLoss.md
@@ -1,3 +1,5 @@
+description: For approximate MRR loss.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.losses.ApproxMRRLoss" />
 <meta itemprop="path" content="Stable" />
@@ -11,10 +13,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L368-L380">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -25,7 +26,8 @@ For approximate MRR loss.
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.losses.ApproxMRRLoss(
-    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None
+    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None,
+    temperature=0.1
 )
 </code></pre>
 
@@ -47,8 +49,8 @@ option will be determined by the usage context. For almost all cases
 this defaults to `SUM_OVER_BATCH_SIZE`. When used with
 `tf.distribute.Strategy`, outside of built-in training loops such as
 `tf.keras` `compile` and `fit`, using `AUTO` or `SUM_OVER_BATCH_SIZE`
-will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training)
-for more details.
+will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training) for
+more details.
 </td>
 </tr><tr>
 <td>
@@ -64,10 +66,13 @@ Optional name for the op.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L301-L308">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_config(
-    config
+    config, custom_objects=None
 )
 </code></pre>
 
@@ -89,7 +94,6 @@ Output of `get_config()`.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -103,6 +107,9 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L291-L299">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_config()
 </code></pre>
@@ -111,7 +118,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L168-L173">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ApproxNDCGLoss.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ApproxNDCGLoss.md
index 589cefd..21373ce 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ApproxNDCGLoss.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ApproxNDCGLoss.md
@@ -1,3 +1,5 @@
+description: For approximate NDCG loss.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.losses.ApproxNDCGLoss" />
 <meta itemprop="path" content="Stable" />
@@ -11,10 +13,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L384-L396">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -25,7 +26,8 @@ For approximate NDCG loss.
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.losses.ApproxNDCGLoss(
-    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None
+    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None,
+    temperature=0.1
 )
 </code></pre>
 
@@ -47,8 +49,8 @@ option will be determined by the usage context. For almost all cases
 this defaults to `SUM_OVER_BATCH_SIZE`. When used with
 `tf.distribute.Strategy`, outside of built-in training loops such as
 `tf.keras` `compile` and `fit`, using `AUTO` or `SUM_OVER_BATCH_SIZE`
-will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training)
-for more details.
+will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training) for
+more details.
 </td>
 </tr><tr>
 <td>
@@ -64,10 +66,13 @@ Optional name for the op.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L301-L308">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_config(
-    config
+    config, custom_objects=None
 )
 </code></pre>
 
@@ -89,7 +94,6 @@ Output of `get_config()`.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -103,6 +107,9 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L291-L299">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_config()
 </code></pre>
@@ -111,7 +118,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L168-L173">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ClickEMLoss.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ClickEMLoss.md
new file mode 100644
index 0000000..9f21018
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ClickEMLoss.md
@@ -0,0 +1,129 @@
+description: For click EM loss.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.losses.ClickEMLoss" />
+<meta itemprop="path" content="Stable" />
+<meta itemprop="property" content="__call__"/>
+<meta itemprop="property" content="__init__"/>
+<meta itemprop="property" content="from_config"/>
+<meta itemprop="property" content="get_config"/>
+</div>
+
+# tfr.keras.losses.ClickEMLoss
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L441-L463">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+For click EM loss.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.losses.ClickEMLoss(
+    reduction=tf.losses.Reduction.AUTO, name=None, exam_loss_weight=1.0,
+    rel_loss_weight=1.0
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`reduction`
+</td>
+<td>
+(Optional) Type of `tf.keras.losses.Reduction` to apply to
+loss. Default value is `AUTO`. `AUTO` indicates that the reduction
+option will be determined by the usage context. For almost all cases
+this defaults to `SUM_OVER_BATCH_SIZE`. When used with
+`tf.distribute.Strategy`, outside of built-in training loops such as
+`tf.keras` `compile` and `fit`, using `AUTO` or `SUM_OVER_BATCH_SIZE`
+will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training) for
+more details.
+</td>
+</tr><tr>
+<td>
+`name`
+</td>
+<td>
+Optional name for the op.
+</td>
+</tr>
+</table>
+
+## Methods
+
+<h3 id="from_config"><code>from_config</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@classmethod</code>
+<code>from_config(
+    config
+)
+</code></pre>
+
+Instantiates a `Loss` from its config (output of `get_config()`).
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`config`
+</td>
+<td>
+Output of `get_config()`.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+A `Loss` instance.
+</td>
+</tr>
+
+</table>
+
+<h3 id="get_config"><code>get_config</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L457-L463">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_config()
+</code></pre>
+
+Returns the config dictionary for a `Loss` instance.
+
+<h3 id="__call__"><code>__call__</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L168-L173">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>__call__(
+    y_true, y_pred, sample_weight=None
+)
+</code></pre>
+
+See tf.keras.losses.Loss.
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/DCGLambdaWeight.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/DCGLambdaWeight.md
new file mode 100644
index 0000000..f6227b8
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/DCGLambdaWeight.md
@@ -0,0 +1,115 @@
+description: Keras serializable class for DCG.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.losses.DCGLambdaWeight" />
+<meta itemprop="path" content="Stable" />
+<meta itemprop="property" content="__init__"/>
+<meta itemprop="property" content="get_config"/>
+<meta itemprop="property" content="individual_weights"/>
+<meta itemprop="property" content="pair_weights"/>
+</div>
+
+# tfr.keras.losses.DCGLambdaWeight
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L85-L107">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Keras serializable class for DCG.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.losses.DCGLambdaWeight(
+    topn=None, gain_fn=None, rank_discount_fn=None, normalized=False,
+    smooth_fraction=0.0, **kwargs
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`topn`
+</td>
+<td>
+(int) The topn for the DCG metric.
+</td>
+</tr><tr>
+<td>
+`gain_fn`
+</td>
+<td>
+(function) Transforms labels.
+</td>
+</tr><tr>
+<td>
+`rank_discount_fn`
+</td>
+<td>
+(function) The rank discount function.
+</td>
+</tr><tr>
+<td>
+`normalized`
+</td>
+<td>
+(bool) If True, normalize weight by the max DCG.
+</td>
+</tr><tr>
+<td>
+`smooth_fraction`
+</td>
+<td>
+(float) parameter to control the contribution from
+LambdaMART.
+</td>
+</tr>
+</table>
+
+## Methods
+
+<h3 id="get_config"><code>get_config</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L100-L107">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_config()
+</code></pre>
+
+<h3 id="individual_weights"><code>individual_weights</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L273-L288">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>individual_weights(
+    labels, ranks
+)
+</code></pre>
+
+See `_LambdaWeight`.
+
+<h3 id="pair_weights"><code>pair_weights</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L213-L271">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>pair_weights(
+    labels, ranks
+)
+</code></pre>
+
+See `_LambdaWeight`.
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/GumbelApproxNDCGLoss.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/GumbelApproxNDCGLoss.md
index 062e946..4991f8c 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/GumbelApproxNDCGLoss.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/GumbelApproxNDCGLoss.md
@@ -1,3 +1,5 @@
+description: For Gumbel approximate NDCG loss.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.losses.GumbelApproxNDCGLoss" />
 <meta itemprop="path" content="Stable" />
@@ -11,10 +13,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L400-L437">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -28,7 +29,7 @@ Inherits From: [`ApproxNDCGLoss`](../../../tfr/keras/losses/ApproxNDCGLoss.md)
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.losses.GumbelApproxNDCGLoss(
     reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None,
-    sample_size=8, temperature=1.0, seed=None
+    temperature=0.1, sample_size=8, gumbel_temperature=1.0, seed=None
 )
 </code></pre>
 
@@ -50,8 +51,8 @@ option will be determined by the usage context. For almost all cases
 this defaults to `SUM_OVER_BATCH_SIZE`. When used with
 `tf.distribute.Strategy`, outside of built-in training loops such as
 `tf.keras` `compile` and `fit`, using `AUTO` or `SUM_OVER_BATCH_SIZE`
-will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training)
-for more details.
+will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training) for
+more details.
 </td>
 </tr><tr>
 <td>
@@ -67,10 +68,13 @@ Optional name for the op.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L301-L308">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_config(
-    config
+    config, custom_objects=None
 )
 </code></pre>
 
@@ -92,7 +96,6 @@ Output of `get_config()`.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -106,6 +109,9 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L422-L429">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_config()
 </code></pre>
@@ -114,7 +120,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L431-L437">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ListMLELambdaWeight.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ListMLELambdaWeight.md
new file mode 100644
index 0000000..feb2aec
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ListMLELambdaWeight.md
@@ -0,0 +1,85 @@
+description: LambdaWeight for ListMLE cost function.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.losses.ListMLELambdaWeight" />
+<meta itemprop="path" content="Stable" />
+<meta itemprop="property" content="__init__"/>
+<meta itemprop="property" content="get_config"/>
+<meta itemprop="property" content="individual_weights"/>
+<meta itemprop="property" content="pair_weights"/>
+</div>
+
+# tfr.keras.losses.ListMLELambdaWeight
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L144-L152">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+LambdaWeight for ListMLE cost function.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.losses.ListMLELambdaWeight(
+    rank_discount_fn, **kwargs
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`rank_discount_fn`
+</td>
+<td>
+(function) The rank discount function.
+</td>
+</tr>
+</table>
+
+## Methods
+
+<h3 id="get_config"><code>get_config</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L149-L152">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_config()
+</code></pre>
+
+<h3 id="individual_weights"><code>individual_weights</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L355-L361">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>individual_weights(
+    labels, ranks
+)
+</code></pre>
+
+See `_LambdaWeight`.
+
+<h3 id="pair_weights"><code>pair_weights</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L351-L353">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>pair_weights(
+    labels, ranks
+)
+</code></pre>
+
+See `_LambdaWeight`.
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ListMLELoss.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ListMLELoss.md
index 22cb9e7..f525bc1 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ListMLELoss.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/ListMLELoss.md
@@ -1,3 +1,5 @@
+description: For List MLE loss.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.losses.ListMLELoss" />
 <meta itemprop="path" content="Stable" />
@@ -11,10 +13,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L352-L364">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -25,7 +26,8 @@ For List MLE loss.
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.losses.ListMLELoss(
-    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None
+    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None,
+    temperature=1.0
 )
 </code></pre>
 
@@ -47,8 +49,8 @@ option will be determined by the usage context. For almost all cases
 this defaults to `SUM_OVER_BATCH_SIZE`. When used with
 `tf.distribute.Strategy`, outside of built-in training loops such as
 `tf.keras` `compile` and `fit`, using `AUTO` or `SUM_OVER_BATCH_SIZE`
-will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training)
-for more details.
+will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training) for
+more details.
 </td>
 </tr><tr>
 <td>
@@ -64,10 +66,13 @@ Optional name for the op.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L301-L308">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_config(
-    config
+    config, custom_objects=None
 )
 </code></pre>
 
@@ -89,7 +94,6 @@ Output of `get_config()`.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -103,6 +107,9 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L291-L299">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_config()
 </code></pre>
@@ -111,7 +118,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L168-L173">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/MeanSquaredLoss.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/MeanSquaredLoss.md
index aa74d2a..865b10f 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/MeanSquaredLoss.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/MeanSquaredLoss.md
@@ -1,3 +1,5 @@
+description: For mean squared error loss.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.losses.MeanSquaredLoss" />
 <meta itemprop="path" content="Stable" />
@@ -11,10 +13,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L477-L483">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -47,8 +48,8 @@ option will be determined by the usage context. For almost all cases
 this defaults to `SUM_OVER_BATCH_SIZE`. When used with
 `tf.distribute.Strategy`, outside of built-in training loops such as
 `tf.keras` `compile` and `fit`, using `AUTO` or `SUM_OVER_BATCH_SIZE`
-will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training)
-for more details.
+will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training) for
+more details.
 </td>
 </tr><tr>
 <td>
@@ -89,7 +90,6 @@ Output of `get_config()`.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -111,7 +111,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L168-L173">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/NDCGLambdaWeight.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/NDCGLambdaWeight.md
new file mode 100644
index 0000000..4957b45
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/NDCGLambdaWeight.md
@@ -0,0 +1,116 @@
+description: Keras serializable class for NDCG.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.losses.NDCGLambdaWeight" />
+<meta itemprop="path" content="Stable" />
+<meta itemprop="property" content="__init__"/>
+<meta itemprop="property" content="get_config"/>
+<meta itemprop="property" content="individual_weights"/>
+<meta itemprop="property" content="pair_weights"/>
+</div>
+
+# tfr.keras.losses.NDCGLambdaWeight
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L111-L125">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Keras serializable class for NDCG.
+
+Inherits From: [`DCGLambdaWeight`](../../../tfr/keras/losses/DCGLambdaWeight.md)
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.losses.NDCGLambdaWeight(
+    topn=None, gain_fn=None, rank_discount_fn=None, smooth_fraction=0.0, **kwargs
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`topn`
+</td>
+<td>
+(int) The topn for the DCG metric.
+</td>
+</tr><tr>
+<td>
+`gain_fn`
+</td>
+<td>
+(function) Transforms labels.
+</td>
+</tr><tr>
+<td>
+`rank_discount_fn`
+</td>
+<td>
+(function) The rank discount function.
+</td>
+</tr><tr>
+<td>
+`normalized`
+</td>
+<td>
+(bool) If True, normalize weight by the max DCG.
+</td>
+</tr><tr>
+<td>
+`smooth_fraction`
+</td>
+<td>
+(float) parameter to control the contribution from
+LambdaMART.
+</td>
+</tr>
+</table>
+
+## Methods
+
+<h3 id="get_config"><code>get_config</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L100-L107">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_config()
+</code></pre>
+
+<h3 id="individual_weights"><code>individual_weights</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L273-L288">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>individual_weights(
+    labels, ranks
+)
+</code></pre>
+
+See `_LambdaWeight`.
+
+<h3 id="pair_weights"><code>pair_weights</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L213-L271">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>pair_weights(
+    labels, ranks
+)
+</code></pre>
+
+See `_LambdaWeight`.
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PairwiseHingeLoss.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PairwiseHingeLoss.md
index 0920ed6..0e4d046 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PairwiseHingeLoss.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PairwiseHingeLoss.md
@@ -1,3 +1,5 @@
+description: For pairwise hinge loss.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.losses.PairwiseHingeLoss" />
 <meta itemprop="path" content="Stable" />
@@ -11,10 +13,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L228-L241">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -25,7 +26,8 @@ For pairwise hinge loss.
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.losses.PairwiseHingeLoss(
-    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None
+    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None,
+    temperature=1.0, **kwargs
 )
 </code></pre>
 
@@ -47,8 +49,8 @@ option will be determined by the usage context. For almost all cases
 this defaults to `SUM_OVER_BATCH_SIZE`. When used with
 `tf.distribute.Strategy`, outside of built-in training loops such as
 `tf.keras` `compile` and `fit`, using `AUTO` or `SUM_OVER_BATCH_SIZE`
-will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training)
-for more details.
+will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training) for
+more details.
 </td>
 </tr><tr>
 <td>
@@ -64,10 +66,13 @@ Optional name for the op.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L206-L213">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_config(
-    config
+    config, custom_objects=None
 )
 </code></pre>
 
@@ -89,7 +94,6 @@ Output of `get_config()`.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -103,6 +107,9 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L196-L204">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_config()
 </code></pre>
@@ -111,7 +118,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L168-L173">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PairwiseLogisticLoss.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PairwiseLogisticLoss.md
index ffb250c..f8d1821 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PairwiseLogisticLoss.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PairwiseLogisticLoss.md
@@ -1,3 +1,5 @@
+description: For pairwise logistic loss.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.losses.PairwiseLogisticLoss" />
 <meta itemprop="path" content="Stable" />
@@ -11,10 +13,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L245-L258">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -25,7 +26,8 @@ For pairwise logistic loss.
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.losses.PairwiseLogisticLoss(
-    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None
+    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None,
+    temperature=1.0, **kwargs
 )
 </code></pre>
 
@@ -47,8 +49,8 @@ option will be determined by the usage context. For almost all cases
 this defaults to `SUM_OVER_BATCH_SIZE`. When used with
 `tf.distribute.Strategy`, outside of built-in training loops such as
 `tf.keras` `compile` and `fit`, using `AUTO` or `SUM_OVER_BATCH_SIZE`
-will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training)
-for more details.
+will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training) for
+more details.
 </td>
 </tr><tr>
 <td>
@@ -64,10 +66,13 @@ Optional name for the op.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L206-L213">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_config(
-    config
+    config, custom_objects=None
 )
 </code></pre>
 
@@ -89,7 +94,6 @@ Output of `get_config()`.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -103,6 +107,9 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L196-L204">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_config()
 </code></pre>
@@ -111,7 +118,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L168-L173">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PairwiseSoftZeroOneLoss.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PairwiseSoftZeroOneLoss.md
index d3673f6..292bc17 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PairwiseSoftZeroOneLoss.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PairwiseSoftZeroOneLoss.md
@@ -1,3 +1,5 @@
+description: For pairwise soft zero-one loss.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.losses.PairwiseSoftZeroOneLoss" />
 <meta itemprop="path" content="Stable" />
@@ -11,10 +13,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L262-L275">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -25,7 +26,8 @@ For pairwise soft zero-one loss.
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.losses.PairwiseSoftZeroOneLoss(
-    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None
+    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None,
+    temperature=1.0, **kwargs
 )
 </code></pre>
 
@@ -47,8 +49,8 @@ option will be determined by the usage context. For almost all cases
 this defaults to `SUM_OVER_BATCH_SIZE`. When used with
 `tf.distribute.Strategy`, outside of built-in training loops such as
 `tf.keras` `compile` and `fit`, using `AUTO` or `SUM_OVER_BATCH_SIZE`
-will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training)
-for more details.
+will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training) for
+more details.
 </td>
 </tr><tr>
 <td>
@@ -64,10 +66,13 @@ Optional name for the op.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L206-L213">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_config(
-    config
+    config, custom_objects=None
 )
 </code></pre>
 
@@ -89,7 +94,6 @@ Output of `get_config()`.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -103,6 +107,9 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L196-L204">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_config()
 </code></pre>
@@ -111,7 +118,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L168-L173">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PrecisionLambdaWeight.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PrecisionLambdaWeight.md
new file mode 100644
index 0000000..c771794
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/PrecisionLambdaWeight.md
@@ -0,0 +1,172 @@
+description: Keras serializable class for Precision.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.losses.PrecisionLambdaWeight" />
+<meta itemprop="path" content="Stable" />
+<meta itemprop="property" content="__init__"/>
+<meta itemprop="property" content="get_config"/>
+<meta itemprop="property" content="individual_weights"/>
+<meta itemprop="property" content="pair_weights"/>
+</div>
+
+# tfr.keras.losses.PrecisionLambdaWeight
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L129-L140">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Keras serializable class for Precision.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.losses.PrecisionLambdaWeight(
+    topn=None, positive_fn=None, **kwargs
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`topn`
+</td>
+<td>
+(int) The K in Precision@K metric.
+</td>
+</tr><tr>
+<td>
+`positive_fn`
+</td>
+<td>
+(function): A function on `Tensor` that output boolean True
+for positive examples. The rest are negative examples.
+</td>
+</tr>
+</table>
+
+## Methods
+
+<h3 id="get_config"><code>get_config</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L136-L140">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_config()
+</code></pre>
+
+<h3 id="individual_weights"><code>individual_weights</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L161-L173">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>individual_weights(
+    labels, ranks
+)
+</code></pre>
+
+Returns the weight `Tensor` for individual examples.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`labels`
+</td>
+<td>
+A dense `Tensor` of labels with shape [batch_size, list_size].
+</td>
+</tr><tr>
+<td>
+`ranks`
+</td>
+<td>
+A dense `Tensor` of ranks with the same shape as `labels` that are
+sorted by logits.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+A `Tensor` that can weight individual examples.
+</td>
+</tr>
+
+</table>
+
+<h3 id="pair_weights"><code>pair_weights</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L307-L335">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>pair_weights(
+    labels, ranks
+)
+</code></pre>
+
+See `_LambdaWeight`.
+
+The current implementation here is that for any pairs of documents i and j, we
+set the weight to be 1 if - i and j have different labels. - i <= topn and j >
+topn or i > topn and j <= topn. This is exactly the same as the original
+LambdaRank method. The weight is the gain of swapping a pair of documents.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`labels`
+</td>
+<td>
+A dense `Tensor` of labels with shape [batch_size, list_size].
+</td>
+</tr><tr>
+<td>
+`ranks`
+</td>
+<td>
+A dense `Tensor` of ranks with the same shape as `labels` that are
+sorted by logits.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+A `Tensor` that can weight example pairs.
+</td>
+</tr>
+
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/RankingLossKey.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/RankingLossKey.md
index 7c46cd4..71ff82b 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/RankingLossKey.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/RankingLossKey.md
@@ -1,3 +1,5 @@
+description: Ranking loss key strings.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.losses.RankingLossKey" />
 <meta itemprop="path" content="Stable" />
@@ -18,10 +20,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L10-L23">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -32,22 +33,88 @@ Ranking loss key strings.
 
 <!-- Placeholder for "Used in" -->
 
-## Class Variables
-
-*   `APPROX_MRR_LOSS = 'approx_mrr_loss'` <a id="APPROX_MRR_LOSS"></a>
-*   `APPROX_NDCG_LOSS = 'approx_ndcg_loss'` <a id="APPROX_NDCG_LOSS"></a>
-*   `GUMBEL_APPROX_NDCG_LOSS = 'gumbel_approx_ndcg_loss'`
-    <a id="GUMBEL_APPROX_NDCG_LOSS"></a>
-*   `LIST_MLE_LOSS = 'list_mle_loss'` <a id="LIST_MLE_LOSS"></a>
-*   `MEAN_SQUARED_LOSS = 'mean_squared_loss'` <a id="MEAN_SQUARED_LOSS"></a>
-*   `PAIRWISE_HINGE_LOSS = 'pairwise_hinge_loss'`
-    <a id="PAIRWISE_HINGE_LOSS"></a>
-*   `PAIRWISE_LOGISTIC_LOSS = 'pairwise_logistic_loss'`
-    <a id="PAIRWISE_LOGISTIC_LOSS"></a>
-*   `PAIRWISE_SOFT_ZERO_ONE_LOSS = 'pairwise_soft_zero_one_loss'`
-    <a id="PAIRWISE_SOFT_ZERO_ONE_LOSS"></a>
-*   `SIGMOID_CROSS_ENTROPY_LOSS = 'sigmoid_cross_entropy_loss'`
-    <a id="SIGMOID_CROSS_ENTROPY_LOSS"></a>
-*   `SOFTMAX_LOSS = 'softmax_loss'` <a id="SOFTMAX_LOSS"></a>
-*   `UNIQUE_SOFTMAX_LOSS = 'unique_softmax_loss'`
-    <a id="UNIQUE_SOFTMAX_LOSS"></a>
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Class Variables</h2></th></tr>
+
+<tr>
+<td>
+APPROX_MRR_LOSS<a id="APPROX_MRR_LOSS"></a>
+</td>
+<td>
+`'approx_mrr_loss'`
+</td>
+</tr><tr>
+<td>
+APPROX_NDCG_LOSS<a id="APPROX_NDCG_LOSS"></a>
+</td>
+<td>
+`'approx_ndcg_loss'`
+</td>
+</tr><tr>
+<td>
+GUMBEL_APPROX_NDCG_LOSS<a id="GUMBEL_APPROX_NDCG_LOSS"></a>
+</td>
+<td>
+`'gumbel_approx_ndcg_loss'`
+</td>
+</tr><tr>
+<td>
+LIST_MLE_LOSS<a id="LIST_MLE_LOSS"></a>
+</td>
+<td>
+`'list_mle_loss'`
+</td>
+</tr><tr>
+<td>
+MEAN_SQUARED_LOSS<a id="MEAN_SQUARED_LOSS"></a>
+</td>
+<td>
+`'mean_squared_loss'`
+</td>
+</tr><tr>
+<td>
+PAIRWISE_HINGE_LOSS<a id="PAIRWISE_HINGE_LOSS"></a>
+</td>
+<td>
+`'pairwise_hinge_loss'`
+</td>
+</tr><tr>
+<td>
+PAIRWISE_LOGISTIC_LOSS<a id="PAIRWISE_LOGISTIC_LOSS"></a>
+</td>
+<td>
+`'pairwise_logistic_loss'`
+</td>
+</tr><tr>
+<td>
+PAIRWISE_SOFT_ZERO_ONE_LOSS<a id="PAIRWISE_SOFT_ZERO_ONE_LOSS"></a>
+</td>
+<td>
+`'pairwise_soft_zero_one_loss'`
+</td>
+</tr><tr>
+<td>
+SIGMOID_CROSS_ENTROPY_LOSS<a id="SIGMOID_CROSS_ENTROPY_LOSS"></a>
+</td>
+<td>
+`'sigmoid_cross_entropy_loss'`
+</td>
+</tr><tr>
+<td>
+SOFTMAX_LOSS<a id="SOFTMAX_LOSS"></a>
+</td>
+<td>
+`'softmax_loss'`
+</td>
+</tr><tr>
+<td>
+UNIQUE_SOFTMAX_LOSS<a id="UNIQUE_SOFTMAX_LOSS"></a>
+</td>
+<td>
+`'unique_softmax_loss'`
+</td>
+</tr>
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/SigmoidCrossEntropyLoss.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/SigmoidCrossEntropyLoss.md
index 82ae497..4f32ca0 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/SigmoidCrossEntropyLoss.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/SigmoidCrossEntropyLoss.md
@@ -1,3 +1,5 @@
+description: For sigmoid cross-entropy loss.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.losses.SigmoidCrossEntropyLoss" />
 <meta itemprop="path" content="Stable" />
@@ -11,10 +13,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L467-L473">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -47,8 +48,8 @@ option will be determined by the usage context. For almost all cases
 this defaults to `SUM_OVER_BATCH_SIZE`. When used with
 `tf.distribute.Strategy`, outside of built-in training loops such as
 `tf.keras` `compile` and `fit`, using `AUTO` or `SUM_OVER_BATCH_SIZE`
-will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training)
-for more details.
+will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training) for
+more details.
 </td>
 </tr><tr>
 <td>
@@ -89,7 +90,6 @@ Output of `get_config()`.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -111,7 +111,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L168-L173">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/SoftmaxLoss.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/SoftmaxLoss.md
index bf8e967..1e029cb 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/SoftmaxLoss.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/SoftmaxLoss.md
@@ -1,3 +1,5 @@
+description: For softmax cross entropy loss.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.losses.SoftmaxLoss" />
 <meta itemprop="path" content="Stable" />
@@ -11,10 +13,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L312-L332">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -25,7 +26,8 @@ For softmax cross entropy loss.
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.losses.SoftmaxLoss(
-    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None
+    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None,
+    temperature=1.0, **kwargs
 )
 </code></pre>
 
@@ -47,8 +49,8 @@ option will be determined by the usage context. For almost all cases
 this defaults to `SUM_OVER_BATCH_SIZE`. When used with
 `tf.distribute.Strategy`, outside of built-in training loops such as
 `tf.keras` `compile` and `fit`, using `AUTO` or `SUM_OVER_BATCH_SIZE`
-will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training)
-for more details.
+will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training) for
+more details.
 </td>
 </tr><tr>
 <td>
@@ -64,10 +66,13 @@ Optional name for the op.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L301-L308">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_config(
-    config
+    config, custom_objects=None
 )
 </code></pre>
 
@@ -89,7 +94,6 @@ Output of `get_config()`.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -103,6 +107,9 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L291-L299">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_config()
 </code></pre>
@@ -111,7 +118,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L327-L332">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/UniqueSoftmaxLoss.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/UniqueSoftmaxLoss.md
index e97b2e4..4b48497 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/UniqueSoftmaxLoss.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/UniqueSoftmaxLoss.md
@@ -1,3 +1,5 @@
+description: For unique softmax cross entropy loss.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.losses.UniqueSoftmaxLoss" />
 <meta itemprop="path" content="Stable" />
@@ -11,10 +13,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L336-L348">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -25,7 +26,8 @@ For unique softmax cross entropy loss.
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.losses.UniqueSoftmaxLoss(
-    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None
+    reduction=tf.losses.Reduction.AUTO, name=None, lambda_weight=None,
+    temperature=1.0
 )
 </code></pre>
 
@@ -47,8 +49,8 @@ option will be determined by the usage context. For almost all cases
 this defaults to `SUM_OVER_BATCH_SIZE`. When used with
 `tf.distribute.Strategy`, outside of built-in training loops such as
 `tf.keras` `compile` and `fit`, using `AUTO` or `SUM_OVER_BATCH_SIZE`
-will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training)
-for more details.
+will raise an error. Please see this custom training [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training) for
+more details.
 </td>
 </tr><tr>
 <td>
@@ -64,10 +66,13 @@ Optional name for the op.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L301-L308">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
 <code>from_config(
-    config
+    config, custom_objects=None
 )
 </code></pre>
 
@@ -102,6 +107,9 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L291-L299">View
+source</a>
+
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_config()
 </code></pre>
@@ -110,7 +118,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L168-L173">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/get.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/get.md
index 4aa8784..48512f3 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/get.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/losses/get.md
@@ -1,3 +1,5 @@
+description: Factory method to get a ranking loss class.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.losses.get" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L26-L81">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics.md
index 0c1a05c..e076e25 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics.md
@@ -1,3 +1,5 @@
+description: Keras metrics in TF-Ranking.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.metrics" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
@@ -48,7 +49,16 @@ Implements PrecisionIA@k (Pre-IA@k).
 [`class PrecisionMetric`](../../tfr/keras/metrics/PrecisionMetric.md):
 Implements precision@k (P@k).
 
+[`class RankingMetricKey`](../../tfr/keras/metrics/RankingMetricKey.md): Ranking
+metric key strings.
+
+[`class RecallMetric`](../../tfr/keras/metrics/RecallMetric.md): Implements
+recall@k.
+
 ## Functions
 
 [`default_keras_metrics(...)`](../../tfr/keras/metrics/default_keras_metrics.md):
 Returns a list of ranking metrics.
+
+[`get(...)`](../../tfr/keras/metrics/get.md): Factory method to get a list of
+ranking metrics.
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/ARPMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/ARPMetric.md
index d989654..9879c98 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/ARPMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/ARPMetric.md
@@ -1,3 +1,5 @@
+description: Implements average relevance position (ARP).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.metrics.ARPMetric" />
 <meta itemprop="path" content="Stable" />
@@ -24,10 +26,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L184-L189">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -38,7 +39,7 @@ Implements average relevance position (ARP).
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.metrics.ARPMetric(
-    name=None, dtype=None, **kwargs
+    name=None, dtype=None, ragged=False, **kwargs
 )
 </code></pre>
 
@@ -50,11 +51,31 @@ Implements average relevance position (ARP).
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `dtype` </td> <td> Dtype used by
-the weights of the layer, set in the constructor. </td> </tr><tr> <td> `dynamic`
-</td> <td> Whether the layer is dynamic (eager-only); set in the constructor.
-</td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
-layer.
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `input` </td> <td> Retrieves the input
+tensor(s) of a layer.
 
 Only applicable if the layer has exactly one input, i.e. if it is connected to
 one incoming layer. </td> </tr><tr> <td> `input_spec` </td> <td> `InputSpec`
@@ -178,6 +199,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -249,9 +277,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -290,11 +319,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -373,9 +402,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -400,9 +430,10 @@ Instance of `TensorShape`, or list of instances of
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -449,9 +480,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -530,9 +562,10 @@ layer from the config dictionary. It does not handle layer connectivity (handled
 by Network), nor weights (handled by `set_weights`).
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -678,9 +711,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -714,7 +748,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L141-L163">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/AlphaDCGMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/AlphaDCGMetric.md
index 332b641..b9aefdf 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/AlphaDCGMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/AlphaDCGMetric.md
@@ -1,3 +1,5 @@
+description: Implements alpha discounted cumulative gain (alphaDCG).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.metrics.AlphaDCGMetric" />
 <meta itemprop="path" content="Stable" />
@@ -24,10 +26,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L351-L407">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -38,13 +39,16 @@ Implements alpha discounted cumulative gain (alphaDCG).
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.metrics.AlphaDCGMetric(
-    topn=None, alpha=0.5, rank_discount_fn=_DEFAULT_RANK_DISCOUNT_FN, seed=None,
-    dtype=None, name='alpha_dcg_metric', **kwargs
+    name=&#x27;alpha_dcg_metric&#x27;, topn=None, alpha=0.5, rank_discount_fn=None,
+    seed=None, dtype=None, ragged=False, **kwargs
 )
 </code></pre>
 
 <!-- Placeholder for "Used in" -->
 
+The `rank_discount_fn` should be keras serializable. Please see the
+`log2_inverse` above examples when defining user customized functions.
+
 <!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
@@ -52,6 +56,13 @@ Implements alpha discounted cumulative gain (alphaDCG).
 
 <tr>
 <td>
+`name`
+</td>
+<td>
+A string used as the name for this metric.
+</td>
+</tr><tr>
+<td>
 `topn`
 </td>
 <td>
@@ -71,8 +82,10 @@ covering a subtopic of the query.
 `rank_discount_fn`
 </td>
 <td>
-A function of rank discounts. Default is set to
-discount = 1 / log2(rank+1).
+A function of rank discounts. Default is set to discount
+= 1 / log2(rank+1). The `rank_discount_fn` should be keras serializable.
+Please see the `log2_inverse` above as an example when defining user
+customized functions.
 </td>
 </tr><tr>
 <td>
@@ -90,10 +103,12 @@ Data type of the metric output. See `tf.keras.metrics.Metric`.
 </td>
 </tr><tr>
 <td>
-`name`
+`ragged`
 </td>
 <td>
-A string used as the name for this metric.
+A bool indicating whether the supplied tensors are ragged. If
+True y_true, y_pred and sample_weight (if providing per-example weights)
+need to be ragged tensors with compatible shapes.
 </td>
 </tr><tr>
 <td>
@@ -111,11 +126,31 @@ Other keyward arguments used in `tf.keras.metrics.Metric`.
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `dtype` </td> <td> Dtype used by
-the weights of the layer, set in the constructor. </td> </tr><tr> <td> `dynamic`
-</td> <td> Whether the layer is dynamic (eager-only); set in the constructor.
-</td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
-layer.
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `input` </td> <td> Retrieves the input
+tensor(s) of a layer.
 
 Only applicable if the layer has exactly one input, i.e. if it is connected to
 one incoming layer. </td> </tr><tr> <td> `input_spec` </td> <td> `InputSpec`
@@ -239,6 +274,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -310,9 +352,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -351,11 +394,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -434,9 +477,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -461,9 +505,10 @@ Instance of `TensorShape`, or list of instances of
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -510,9 +555,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -591,9 +637,10 @@ layer from the config dictionary. It does not handle layer connectivity (handled
 by Network), nor weights (handled by `set_weights`).
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -620,7 +667,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L399-L407">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -742,9 +789,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -778,7 +826,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L141-L163">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/DCGMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/DCGMetric.md
index 0fff7cf..ecdc0ae 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/DCGMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/DCGMetric.md
@@ -1,3 +1,5 @@
+description: Implements discounted cumulative gain (DCG).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.metrics.DCGMetric" />
 <meta itemprop="path" content="Stable" />
@@ -24,10 +26,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L312-L347">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -38,24 +39,48 @@ Implements discounted cumulative gain (DCG).
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.metrics.DCGMetric(
-    name=None, topn=None, gain_fn=_DEFAULT_GAIN_FN,
-    rank_discount_fn=_DEFAULT_RANK_DISCOUNT_FN, dtype=None, **kwargs
+    name=None, topn=None, gain_fn=None, rank_discount_fn=None, dtype=None,
+    ragged=False, **kwargs
 )
 </code></pre>
 
 <!-- Placeholder for "Used in" -->
 
+The `gain_fn` and `rank_discount_fn` should be keras serializable. Please see
+the `pow_minus_1` and `log2_inverse` above as examples when defining user
+customized functions.
+
 <!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `dtype` </td> <td> Dtype used by
-the weights of the layer, set in the constructor. </td> </tr><tr> <td> `dynamic`
-</td> <td> Whether the layer is dynamic (eager-only); set in the constructor.
-</td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
-layer.
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `input` </td> <td> Retrieves the input
+tensor(s) of a layer.
 
 Only applicable if the layer has exactly one input, i.e. if it is connected to
 one incoming layer. </td> </tr><tr> <td> `input_spec` </td> <td> `InputSpec`
@@ -179,6 +204,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -250,9 +282,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -291,11 +324,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -374,9 +407,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -401,9 +435,10 @@ Instance of `TensorShape`, or list of instances of
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -450,9 +485,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -531,9 +567,10 @@ layer from the config dictionary. It does not handle layer connectivity (handled
 by Network), nor weights (handled by `set_weights`).
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -560,7 +597,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L339-L347">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -682,9 +719,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -718,7 +756,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L141-L163">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MRRMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MRRMetric.md
index 60d0efb..37959d2 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MRRMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MRRMetric.md
@@ -1,3 +1,5 @@
+description: Implements mean reciprocal rank (MRR).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.metrics.MRRMetric" />
 <meta itemprop="path" content="Stable" />
@@ -24,10 +26,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L167-L180">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -38,7 +39,7 @@ Implements mean reciprocal rank (MRR).
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.metrics.MRRMetric(
-    name=None, topn=None, dtype=None, **kwargs
+    name=None, topn=None, dtype=None, ragged=False, **kwargs
 )
 </code></pre>
 
@@ -50,11 +51,31 @@ Implements mean reciprocal rank (MRR).
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `dtype` </td> <td> Dtype used by
-the weights of the layer, set in the constructor. </td> </tr><tr> <td> `dynamic`
-</td> <td> Whether the layer is dynamic (eager-only); set in the constructor.
-</td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
-layer.
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `input` </td> <td> Retrieves the input
+tensor(s) of a layer.
 
 Only applicable if the layer has exactly one input, i.e. if it is connected to
 one incoming layer. </td> </tr><tr> <td> `input_spec` </td> <td> `InputSpec`
@@ -178,6 +199,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -249,9 +277,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -290,11 +319,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -373,9 +402,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -400,9 +430,10 @@ Instance of `TensorShape`, or list of instances of
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -449,9 +480,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -530,9 +562,10 @@ layer from the config dictionary. It does not handle layer connectivity (handled
 by Network), nor weights (handled by `set_weights`).
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -559,7 +592,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L175-L180">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -681,9 +714,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -717,7 +751,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L141-L163">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MeanAveragePrecisionMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MeanAveragePrecisionMetric.md
index 6fb5fa1..c627ee1 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MeanAveragePrecisionMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/MeanAveragePrecisionMetric.md
@@ -1,3 +1,5 @@
+description: Implements mean average precision (MAP).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.metrics.MeanAveragePrecisionMetric" />
 <meta itemprop="path" content="Stable" />
@@ -24,10 +26,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L253-L269">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -38,7 +39,7 @@ Implements mean average precision (MAP).
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.metrics.MeanAveragePrecisionMetric(
-    name=None, topn=None, dtype=None, **kwargs
+    name=None, topn=None, dtype=None, ragged=False, **kwargs
 )
 </code></pre>
 
@@ -50,11 +51,31 @@ Implements mean average precision (MAP).
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `dtype` </td> <td> Dtype used by
-the weights of the layer, set in the constructor. </td> </tr><tr> <td> `dynamic`
-</td> <td> Whether the layer is dynamic (eager-only); set in the constructor.
-</td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
-layer.
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `input` </td> <td> Retrieves the input
+tensor(s) of a layer.
 
 Only applicable if the layer has exactly one input, i.e. if it is connected to
 one incoming layer. </td> </tr><tr> <td> `input_spec` </td> <td> `InputSpec`
@@ -178,6 +199,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -249,9 +277,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -290,11 +319,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -373,9 +402,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -400,9 +430,10 @@ Instance of `TensorShape`, or list of instances of
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -449,9 +480,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -530,9 +562,10 @@ layer from the config dictionary. It does not handle layer connectivity (handled
 by Network), nor weights (handled by `set_weights`).
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -559,7 +592,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L263-L269">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -681,9 +714,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -717,7 +751,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L141-L163">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/NDCGMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/NDCGMetric.md
index 390269a..97b4a70 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/NDCGMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/NDCGMetric.md
@@ -1,3 +1,5 @@
+description: Implements normalized discounted cumulative gain (NDCG).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.metrics.NDCGMetric" />
 <meta itemprop="path" content="Stable" />
@@ -24,10 +26,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L273-L308">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -38,24 +39,48 @@ Implements normalized discounted cumulative gain (NDCG).
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.metrics.NDCGMetric(
-    name=None, topn=None, gain_fn=_DEFAULT_GAIN_FN,
-    rank_discount_fn=_DEFAULT_RANK_DISCOUNT_FN, dtype=None, **kwargs
+    name=None, topn=None, gain_fn=None, rank_discount_fn=None, dtype=None,
+    ragged=False, **kwargs
 )
 </code></pre>
 
 <!-- Placeholder for "Used in" -->
 
+The `gain_fn` and `rank_discount_fn` should be keras serializable. Please see
+the `pow_minus_1` and `log2_inverse` above as examples when defining user
+customized functions.
+
 <!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `dtype` </td> <td> Dtype used by
-the weights of the layer, set in the constructor. </td> </tr><tr> <td> `dynamic`
-</td> <td> Whether the layer is dynamic (eager-only); set in the constructor.
-</td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
-layer.
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `input` </td> <td> Retrieves the input
+tensor(s) of a layer.
 
 Only applicable if the layer has exactly one input, i.e. if it is connected to
 one incoming layer. </td> </tr><tr> <td> `input_spec` </td> <td> `InputSpec`
@@ -179,6 +204,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -250,9 +282,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -291,11 +324,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -374,9 +407,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -401,9 +435,10 @@ Instance of `TensorShape`, or list of instances of
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -450,9 +485,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -531,9 +567,10 @@ layer from the config dictionary. It does not handle layer connectivity (handled
 by Network), nor weights (handled by `set_weights`).
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -560,7 +597,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L300-L308">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -682,9 +719,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -718,7 +756,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L141-L163">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/OPAMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/OPAMetric.md
index 26a54ff..a7baa5f 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/OPAMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/OPAMetric.md
@@ -1,3 +1,5 @@
+description: Implements ordered pair accuracy (OPA).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.metrics.OPAMetric" />
 <meta itemprop="path" content="Stable" />
@@ -24,10 +26,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L411-L416">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -38,7 +39,7 @@ Implements ordered pair accuracy (OPA).
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.metrics.OPAMetric(
-    name=None, dtype=None, **kwargs
+    name=None, dtype=None, ragged=False, **kwargs
 )
 </code></pre>
 
@@ -50,11 +51,31 @@ Implements ordered pair accuracy (OPA).
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `dtype` </td> <td> Dtype used by
-the weights of the layer, set in the constructor. </td> </tr><tr> <td> `dynamic`
-</td> <td> Whether the layer is dynamic (eager-only); set in the constructor.
-</td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
-layer.
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `input` </td> <td> Retrieves the input
+tensor(s) of a layer.
 
 Only applicable if the layer has exactly one input, i.e. if it is connected to
 one incoming layer. </td> </tr><tr> <td> `input_spec` </td> <td> `InputSpec`
@@ -178,6 +199,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -249,9 +277,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -290,11 +319,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -373,9 +402,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -400,9 +430,10 @@ Instance of `TensorShape`, or list of instances of
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -449,9 +480,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -530,9 +562,10 @@ layer from the config dictionary. It does not handle layer connectivity (handled
 by Network), nor weights (handled by `set_weights`).
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -678,9 +711,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -714,7 +748,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L141-L163">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionIAMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionIAMetric.md
index ffc6376..395e244 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionIAMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionIAMetric.md
@@ -1,3 +1,5 @@
+description: Implements PrecisionIA@k (Pre-IA@k).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.metrics.PrecisionIAMetric" />
 <meta itemprop="path" content="Stable" />
@@ -24,10 +26,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L230-L249">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -38,7 +39,7 @@ Implements PrecisionIA@k (Pre-IA@k).
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.metrics.PrecisionIAMetric(
-    topn=None, dtype=None, name='precision_ia_metric', **kwargs
+    name=None, topn=None, dtype=None, ragged=False, **kwargs
 )
 </code></pre>
 
@@ -50,11 +51,31 @@ Implements PrecisionIA@k (Pre-IA@k).
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `dtype` </td> <td> Dtype used by
-the weights of the layer, set in the constructor. </td> </tr><tr> <td> `dynamic`
-</td> <td> Whether the layer is dynamic (eager-only); set in the constructor.
-</td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
-layer.
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `input` </td> <td> Retrieves the input
+tensor(s) of a layer.
 
 Only applicable if the layer has exactly one input, i.e. if it is connected to
 one incoming layer. </td> </tr><tr> <td> `input_spec` </td> <td> `InputSpec`
@@ -178,6 +199,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -249,9 +277,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -290,11 +319,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -373,9 +402,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -400,9 +430,10 @@ Instance of `TensorShape`, or list of instances of
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -449,9 +480,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -530,9 +562,10 @@ layer from the config dictionary. It does not handle layer connectivity (handled
 by Network), nor weights (handled by `set_weights`).
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -559,7 +592,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L244-L249">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -681,9 +714,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -717,7 +751,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L141-L163">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionMetric.md
index 381f147..db4316d 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionMetric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/PrecisionMetric.md
@@ -1,3 +1,5 @@
+description: Implements precision@k (P@k).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.metrics.PrecisionMetric" />
 <meta itemprop="path" content="Stable" />
@@ -24,10 +26,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L193-L207">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -38,7 +39,7 @@ Implements precision@k (P@k).
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.metrics.PrecisionMetric(
-    name=None, topn=None, dtype=None, **kwargs
+    name=None, topn=None, dtype=None, ragged=False, **kwargs
 )
 </code></pre>
 
@@ -50,11 +51,31 @@ Implements precision@k (P@k).
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `dtype` </td> <td> Dtype used by
-the weights of the layer, set in the constructor. </td> </tr><tr> <td> `dynamic`
-</td> <td> Whether the layer is dynamic (eager-only); set in the constructor.
-</td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
-layer.
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `input` </td> <td> Retrieves the input
+tensor(s) of a layer.
 
 Only applicable if the layer has exactly one input, i.e. if it is connected to
 one incoming layer. </td> </tr><tr> <td> `input_spec` </td> <td> `InputSpec`
@@ -178,6 +199,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -249,9 +277,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -290,11 +319,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -373,9 +402,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -400,9 +430,10 @@ Instance of `TensorShape`, or list of instances of
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -449,9 +480,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -530,9 +562,10 @@ layer from the config dictionary. It does not handle layer connectivity (handled
 by Network), nor weights (handled by `set_weights`).
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -559,7 +592,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L202-L207">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -681,9 +714,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -717,7 +751,7 @@ layer's specifications.
 
 <h3 id="update_state"><code>update_state</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L141-L163">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/RankingMetricKey.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/RankingMetricKey.md
new file mode 100644
index 0000000..d47ffd7
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/RankingMetricKey.md
@@ -0,0 +1,103 @@
+description: Ranking metric key strings.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.metrics.RankingMetricKey" />
+<meta itemprop="path" content="Stable" />
+<meta itemprop="property" content="ALPHA_DCG"/>
+<meta itemprop="property" content="ARP"/>
+<meta itemprop="property" content="DCG"/>
+<meta itemprop="property" content="MAP"/>
+<meta itemprop="property" content="MRR"/>
+<meta itemprop="property" content="NDCG"/>
+<meta itemprop="property" content="ORDERED_PAIR_ACCURACY"/>
+<meta itemprop="property" content="PRECISION"/>
+<meta itemprop="property" content="PRECISION_IA"/>
+</div>
+
+# tfr.keras.metrics.RankingMetricKey
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L15-L42">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Ranking metric key strings.
+
+<!-- Placeholder for "Used in" -->
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Class Variables</h2></th></tr>
+
+<tr>
+<td>
+ALPHA_DCG<a id="ALPHA_DCG"></a>
+</td>
+<td>
+`'alpha_dcg'`
+</td>
+</tr><tr>
+<td>
+ARP<a id="ARP"></a>
+</td>
+<td>
+`'arp'`
+</td>
+</tr><tr>
+<td>
+DCG<a id="DCG"></a>
+</td>
+<td>
+`'dcg'`
+</td>
+</tr><tr>
+<td>
+MAP<a id="MAP"></a>
+</td>
+<td>
+`'map'`
+</td>
+</tr><tr>
+<td>
+MRR<a id="MRR"></a>
+</td>
+<td>
+`'mrr'`
+</td>
+</tr><tr>
+<td>
+NDCG<a id="NDCG"></a>
+</td>
+<td>
+`'ndcg'`
+</td>
+</tr><tr>
+<td>
+ORDERED_PAIR_ACCURACY<a id="ORDERED_PAIR_ACCURACY"></a>
+</td>
+<td>
+`'ordered_pair_accuracy'`
+</td>
+</tr><tr>
+<td>
+PRECISION<a id="PRECISION"></a>
+</td>
+<td>
+`'precision'`
+</td>
+</tr><tr>
+<td>
+PRECISION_IA<a id="PRECISION_IA"></a>
+</td>
+<td>
+`'precision_ia'`
+</td>
+</tr>
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/RecallMetric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/RecallMetric.md
new file mode 100644
index 0000000..d702f38
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/RecallMetric.md
@@ -0,0 +1,922 @@
+description: Implements recall@k.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.metrics.RecallMetric" />
+<meta itemprop="path" content="Stable" />
+<meta itemprop="property" content="__call__"/>
+<meta itemprop="property" content="__init__"/>
+<meta itemprop="property" content="__new__"/>
+<meta itemprop="property" content="add_loss"/>
+<meta itemprop="property" content="add_metric"/>
+<meta itemprop="property" content="build"/>
+<meta itemprop="property" content="compute_mask"/>
+<meta itemprop="property" content="compute_output_shape"/>
+<meta itemprop="property" content="count_params"/>
+<meta itemprop="property" content="from_config"/>
+<meta itemprop="property" content="get_config"/>
+<meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="reset_states"/>
+<meta itemprop="property" content="result"/>
+<meta itemprop="property" content="set_weights"/>
+<meta itemprop="property" content="update_state"/>
+<meta itemprop="property" content="with_name_scope"/>
+</div>
+
+# tfr.keras.metrics.RecallMetric
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L212-L226">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Implements recall@k.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.metrics.RecallMetric(
+    name=None, topn=None, dtype=None, ragged=False, **kwargs
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
+
+<tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`dtype` </td> <td> The dtype of the layer weights.
+
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `input` </td> <td> Retrieves the input
+tensor(s) of a layer.
+
+Only applicable if the layer has exactly one input, i.e. if it is connected to
+one incoming layer. </td> </tr><tr> <td> `input_spec` </td> <td> `InputSpec`
+instance(s) describing the input format for this layer.
+
+When you create a layer subclass, you can set `self.input_spec` to enable the
+layer to run input compatibility checks when it is called. Consider a `Conv2D`
+layer: it can only be called on a single input tensor of rank 4. As such, you
+can set, in `__init__()`:
+
+```python
+self.input_spec = tf.keras.layers.InputSpec(ndim=4)
+```
+
+Now, if you try to call the layer on an input that isn't rank 4 (for instance,
+an input of shape `(2,)`, it will raise a nicely-formatted error:
+
+```
+ValueError: Input 0 of layer conv2d is incompatible with the layer:
+expected ndim=4, found ndim=1. Full shape received: [2]
+```
+
+Input checks that can be specified via `input_spec` include: - Structure (e.g. a
+single input, a list of 2 inputs, etc) - Shape - Rank (ndim) - Dtype
+
+For more information, see `tf.keras.layers.InputSpec`. </td> </tr><tr> <td>
+`losses` </td> <td> List of losses added using the `add_loss()` API.
+
+Variable regularization tensors are created when this property is accessed, so
+it is eager safe: accessing `losses` under a `tf.GradientTape` will propagate
+gradients back to the corresponding variables.
+
+```
+>>> class MyLayer(tf.keras.layers.Layer):
+...   def call(self, inputs):
+...     self.add_loss(tf.abs(tf.reduce_mean(inputs)))
+...     return inputs
+>>> l = MyLayer()
+>>> l(np.ones((10, 1)))
+>>> l.losses
+[1.0]
+```
+
+```
+>>> inputs = tf.keras.Input(shape=(10,))
+>>> x = tf.keras.layers.Dense(10)(inputs)
+>>> outputs = tf.keras.layers.Dense(1)(x)
+>>> model = tf.keras.Model(inputs, outputs)
+>>> # Activity regularization.
+>>> len(model.losses)
+0
+>>> model.add_loss(tf.abs(tf.reduce_mean(x)))
+>>> len(model.losses)
+1
+```
+
+```
+>>> inputs = tf.keras.Input(shape=(10,))
+>>> d = tf.keras.layers.Dense(10, kernel_initializer='ones')
+>>> x = d(inputs)
+>>> outputs = tf.keras.layers.Dense(1)(x)
+>>> model = tf.keras.Model(inputs, outputs)
+>>> # Weight regularization.
+>>> model.add_loss(lambda: tf.reduce_mean(d.kernel))
+>>> model.losses
+[<tf.Tensor: shape=(), dtype=float32, numpy=1.0>]
+```
+
+</td> </tr><tr> <td> `metrics` </td> <td> List of metrics added using the
+`add_metric()` API.
+
+```
+>>> input = tf.keras.layers.Input(shape=(3,))
+>>> d = tf.keras.layers.Dense(2)
+>>> output = d(input)
+>>> d.add_metric(tf.reduce_max(output), name='max')
+>>> d.add_metric(tf.reduce_min(output), name='min')
+>>> [m.name for m in d.metrics]
+['max', 'min']
+```
+
+</td> </tr><tr> <td> `name` </td> <td> Name of the layer (string), set in the
+constructor. </td> </tr><tr> <td> `name_scope` </td> <td> Returns a
+`tf.name_scope` instance for this class. </td> </tr><tr> <td>
+`non_trainable_weights` </td> <td> List of all non-trainable weights tracked by
+this layer.
+
+Non-trainable weights are *not* updated during training. They are expected to be
+updated manually in `call()`. </td> </tr><tr> <td> `output` </td> <td> Retrieves
+the output tensor(s) of a layer.
+
+Only applicable if the layer has exactly one output, i.e. if it is connected to
+one incoming layer. </td> </tr><tr> <td> `submodules` </td> <td> Sequence of all
+sub-modules.
+
+Submodules are modules which are properties of this module, or found as
+properties of modules which are properties of this module (and so on).
+
+```
+>>> a = tf.Module()
+>>> b = tf.Module()
+>>> c = tf.Module()
+>>> a.b = b
+>>> b.c = c
+>>> list(a.submodules) == [b, c]
+True
+>>> list(b.submodules) == [c]
+True
+>>> list(c.submodules) == []
+True
+```
+
+</td> </tr><tr> <td> `supports_masking` </td> <td> Whether this layer supports
+computing a mask using `compute_mask`. </td> </tr><tr> <td> `trainable` </td>
+<td>
+
+</td> </tr><tr> <td> `trainable_weights` </td> <td> List of all trainable
+weights tracked by this layer.
+
+Trainable weights are updated via gradient descent during training.
+</td>
+</tr><tr>
+<td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
+`weights`
+</td>
+<td>
+Returns the list of all layer variables/weights.
+</td>
+</tr>
+</table>
+
+## Methods
+
+<h3 id="add_loss"><code>add_loss</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>add_loss(
+    losses, **kwargs
+)
+</code></pre>
+
+Add loss tensor(s), potentially dependent on layer inputs.
+
+Some losses (for instance, activity regularization losses) may be dependent on
+the inputs passed when calling a layer. Hence, when reusing the same layer on
+different inputs `a` and `b`, some entries in `layer.losses` may be dependent on
+`a` and some on `b`. This method automatically keeps track of dependencies.
+
+This method can be used inside a subclassed layer or model's `call` function, in
+which case `losses` should be a Tensor or list of Tensors.
+
+#### Example:
+
+```python
+class MyLayer(tf.keras.layers.Layer):
+  def call(self, inputs):
+    self.add_loss(tf.abs(tf.reduce_mean(inputs)))
+    return inputs
+```
+
+This method can also be called directly on a Functional Model during
+construction. In this case, any loss Tensors passed to this Model must be
+symbolic and be able to be traced back to the model's `Input`s. These losses
+become part of the model's topology and are tracked in `get_config`.
+
+#### Example:
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+# Activity regularization.
+model.add_loss(tf.abs(tf.reduce_mean(x)))
+```
+
+If this is not the case for your loss (if, for example, your loss references a
+`Variable` of one of the model's layers), you can wrap your loss in a
+zero-argument lambda. These losses are not tracked as part of the model's
+topology since they can't be serialized.
+
+#### Example:
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+d = tf.keras.layers.Dense(10)
+x = d(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+# Weight regularization.
+model.add_loss(lambda: tf.reduce_mean(d.kernel))
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`losses`
+</td>
+<td>
+Loss tensor, or list/tuple of tensors. Rather than tensors, losses
+may also be zero-argument callables which create a loss tensor.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Additional keyword arguments for backward compatibility.
+Accepted values:
+inputs - Deprecated, will be automatically inferred.
+</td>
+</tr>
+</table>
+
+<h3 id="add_metric"><code>add_metric</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>add_metric(
+    value, name=None, **kwargs
+)
+</code></pre>
+
+Adds metric tensor to the layer.
+
+This method can be used inside the `call()` method of a subclassed layer or
+model.
+
+```python
+class MyMetricLayer(tf.keras.layers.Layer):
+  def __init__(self):
+    super(MyMetricLayer, self).__init__(name='my_metric_layer')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
+
+  def call(self, inputs):
+    self.add_metric(self.mean(x))
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
+    return inputs
+```
+
+This method can also be called directly on a Functional Model during
+construction. In this case, any tensor passed to this Model must be symbolic and
+be able to be traced back to the model's `Input`s. These metrics become part of
+the model's topology and are tracked when you save the model via `save()`.
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+model.add_metric(math_ops.reduce_sum(x), name='metric_1')
+```
+
+Note: Calling `add_metric()` with the result of a metric object on a Functional
+Model, as shown in the example below, is not supported. This is because we
+cannot trace the metric result tensor back to the model's inputs.
+
+```python
+inputs = tf.keras.Input(shape=(10,))
+x = tf.keras.layers.Dense(10)(inputs)
+outputs = tf.keras.layers.Dense(1)(x)
+model = tf.keras.Model(inputs, outputs)
+model.add_metric(tf.keras.metrics.Mean()(x), name='metric_1')
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`value`
+</td>
+<td>
+Metric tensor.
+</td>
+</tr><tr>
+<td>
+`name`
+</td>
+<td>
+String metric name.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Additional keyword arguments for backward compatibility.
+Accepted values:
+`aggregation` - When the `value` tensor provided is not the result of
+calling a `keras.Metric` instance, it will be aggregated by default
+using a `keras.Metric.Mean`.
+</td>
+</tr>
+</table>
+
+<h3 id="build"><code>build</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>build(
+    input_shape
+)
+</code></pre>
+
+Creates the variables of the layer (optional, for subclass implementers).
+
+This is a method that implementers of subclasses of `Layer` or `Model` can
+override if they need a state-creation step in-between layer instantiation and
+layer call.
+
+This is typically used to create the weights of `Layer` subclasses.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`input_shape`
+</td>
+<td>
+Instance of `TensorShape`, or list of instances of
+`TensorShape` if the layer expects a list of inputs
+(one instance per input).
+</td>
+</tr>
+</table>
+
+<h3 id="compute_mask"><code>compute_mask</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>compute_mask(
+    inputs, mask=None
+)
+</code></pre>
+
+Computes an output mask tensor.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`inputs`
+</td>
+<td>
+Tensor or list of tensors.
+</td>
+</tr><tr>
+<td>
+`mask`
+</td>
+<td>
+Tensor or list of tensors.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+None or a tensor (or list of tensors,
+one per output tensor of the layer).
+</td>
+</tr>
+
+</table>
+
+<h3 id="compute_output_shape"><code>compute_output_shape</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>compute_output_shape(
+    input_shape
+)
+</code></pre>
+
+Computes the output shape of the layer.
+
+If the layer has not been built, this method will call `build` on the layer.
+This assumes that the layer will later be used with inputs that match the input
+shape provided here.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`input_shape`
+</td>
+<td>
+Shape tuple (tuple of integers)
+or list of shape tuples (one per output tensor of the layer).
+Shape tuples can include None for free dimensions,
+instead of an integer.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+An input shape tuple.
+</td>
+</tr>
+
+</table>
+
+<h3 id="count_params"><code>count_params</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>count_params()
+</code></pre>
+
+Count the total number of scalars composing the weights.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+An integer count.
+</td>
+</tr>
+
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Raises</th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+if the layer isn't yet built
+(in which case its weights aren't yet defined).
+</td>
+</tr>
+</table>
+
+<h3 id="from_config"><code>from_config</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@classmethod</code>
+<code>from_config(
+    config
+)
+</code></pre>
+
+Creates a layer from its config.
+
+This method is the reverse of `get_config`, capable of instantiating the same
+layer from the config dictionary. It does not handle layer connectivity (handled
+by Network), nor weights (handled by `set_weights`).
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`config`
+</td>
+<td>
+A Python dictionary, typically the
+output of get_config.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+A layer instance.
+</td>
+</tr>
+
+</table>
+
+<h3 id="get_config"><code>get_config</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L221-L226">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_config()
+</code></pre>
+
+Returns the serializable config of the metric.
+
+<h3 id="get_weights"><code>get_weights</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_weights()
+</code></pre>
+
+Returns the current weights of the layer.
+
+The weights of a layer represent the state of the layer. This function returns
+both trainable and non-trainable weight values associated with this layer as a
+list of Numpy arrays, which can in turn be used to load state into similarly
+parameterized layers.
+
+For example, a Dense layer returns a list of two values-- per-output weights and
+the bias value. These can be used to set the weights of another Dense layer:
+
+```
+>>> a = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(1.))
+>>> a_out = a(tf.convert_to_tensor([[1., 2., 3.]]))
+>>> a.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+>>> b = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(2.))
+>>> b_out = b(tf.convert_to_tensor([[10., 20., 30.]]))
+>>> b.get_weights()
+[array([[2.],
+       [2.],
+       [2.]], dtype=float32), array([0.], dtype=float32)]
+>>> b.set_weights(a.get_weights())
+>>> b.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+Weights values as a list of numpy arrays.
+</td>
+</tr>
+
+</table>
+
+<h3 id="reset_states"><code>reset_states</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>reset_states()
+</code></pre>
+
+Resets all of the metric state variables.
+
+This function is called between epochs/steps, when a metric is evaluated during
+training.
+
+<h3 id="result"><code>result</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>result()
+</code></pre>
+
+Computes and returns the metric value tensor.
+
+Result computation is an idempotent operation that simply calculates the metric
+value using the state variables.
+
+<h3 id="set_weights"><code>set_weights</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>set_weights(
+    weights
+)
+</code></pre>
+
+Sets the weights of the layer, from Numpy arrays.
+
+The weights of a layer represent the state of the layer. This function sets the
+weight values from numpy arrays. The weight values should be passed in the order
+they are created by the layer. Note that the layer's weights must be
+instantiated before calling this function by calling the layer.
+
+For example, a Dense layer returns a list of two values-- per-output weights and
+the bias value. These can be used to set the weights of another Dense layer:
+
+```
+>>> a = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(1.))
+>>> a_out = a(tf.convert_to_tensor([[1., 2., 3.]]))
+>>> a.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+>>> b = tf.keras.layers.Dense(1,
+...   kernel_initializer=tf.constant_initializer(2.))
+>>> b_out = b(tf.convert_to_tensor([[10., 20., 30.]]))
+>>> b.get_weights()
+[array([[2.],
+       [2.],
+       [2.]], dtype=float32), array([0.], dtype=float32)]
+>>> b.set_weights(a.get_weights())
+>>> b.get_weights()
+[array([[1.],
+       [1.],
+       [1.]], dtype=float32), array([0.], dtype=float32)]
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`weights`
+</td>
+<td>
+a list of Numpy arrays. The number
+of arrays and their shape must match
+number of the dimensions of the weights
+of the layer (i.e. it should match the
+output of `get_weights`).
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Raises</th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+If the provided weights list does not match the
+layer's specifications.
+</td>
+</tr>
+</table>
+
+<h3 id="update_state"><code>update_state</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L141-L163">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>update_state(
+    y_true, y_pred, sample_weight=None
+)
+</code></pre>
+
+Accumulates metric statistics.
+
+`y_true` and `y_pred` should have the same shape.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`y_true`
+</td>
+<td>
+The ground truth values.
+</td>
+</tr><tr>
+<td>
+`y_pred`
+</td>
+<td>
+The predicted values.
+</td>
+</tr><tr>
+<td>
+`sample_weight`
+</td>
+<td>
+Optional weighting of each example. Defaults to 1. Can be a
+`Tensor` whose rank is either 0, or the same rank as `y_true`, and must
+be broadcastable to `y_true`.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+Update op.
+</td>
+</tr>
+
+</table>
+
+<h3 id="with_name_scope"><code>with_name_scope</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@classmethod</code>
+<code>with_name_scope(
+    method
+)
+</code></pre>
+
+Decorator to automatically enter the module name scope.
+
+```
+>>> class MyModule(tf.Module):
+...   @tf.Module.with_name_scope
+...   def __call__(self, x):
+...     if not hasattr(self, 'w'):
+...       self.w = tf.Variable(tf.random.normal([x.shape[1], 3]))
+...     return tf.matmul(x, self.w)
+```
+
+Using the above module would produce `tf.Variable`s and `tf.Tensor`s whose names
+included the module name:
+
+```
+>>> mod = MyModule()
+>>> mod(tf.ones([1, 2]))
+<tf.Tensor: shape=(1, 3), dtype=float32, numpy=..., dtype=float32)>
+>>> mod.w
+<tf.Variable 'my_module/Variable:0' shape=(2, 3) dtype=float32,
+numpy=..., dtype=float32)>
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`method`
+</td>
+<td>
+The method to wrap.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+The original method wrapped such that it enters the module's name scope.
+</td>
+</tr>
+
+</table>
+
+<h3 id="__call__"><code>__call__</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>__call__(
+    *args, **kwargs
+)
+</code></pre>
+
+Accumulates statistics and then computes metric result value.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr> <td> `*args` </td> <td>
+
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+A mini-batch of inputs to the Metric,
+passed on to `update_state()`.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+The metric value tensor.
+</td>
+</tr>
+
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/default_keras_metrics.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/default_keras_metrics.md
index 050a4d9..26f414d 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/default_keras_metrics.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/default_keras_metrics.md
@@ -1,3 +1,5 @@
+description: Returns a list of ranking metrics.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.metrics.default_keras_metrics" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L106-L124">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -20,7 +21,7 @@
 Returns a list of ranking metrics.
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
-<code>tfr.keras.metrics.default_keras_metrics()
+<code>tfr.keras.metrics.default_keras_metrics() -> List[tf.keras.metrics.Metric]
 </code></pre>
 
 <!-- Placeholder for "Used in" -->
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/get.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/get.md
new file mode 100644
index 0000000..df1ff79
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/metrics/get.md
@@ -0,0 +1,117 @@
+description: Factory method to get a list of ranking metrics.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.metrics.get" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# tfr.keras.metrics.get
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L45-L103">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Factory method to get a list of ranking metrics.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.metrics.get(
+    key: str,
+    name: Optional[str] = None,
+    dtype: Optional[tf.dtypes.DType] = None,
+    topn: Optional[int] = None,
+    **kwargs
+) -> tf.keras.metrics.Metric
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+
+#### Example Usage:
+
+```python
+  metric = tfr.keras.metics.get(tfr.keras.metrics.RankingMetricKey.MRR)
+```
+
+to get Mean Reciprocal Rank. `python metric =
+tfr.keras.metics.get(tfr.keras.metrics.RankingMetricKey.MRR, topn=2)` to get
+MRR@2.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`key`
+</td>
+<td>
+An attribute of `RankingMetricKey`, defining which metric objects to
+return.
+</td>
+</tr><tr>
+<td>
+`name`
+</td>
+<td>
+Name of metrics.
+</td>
+</tr><tr>
+<td>
+`dtype`
+</td>
+<td>
+Dtype of the metrics.
+</td>
+</tr><tr>
+<td>
+`topn`
+</td>
+<td>
+Cutoff of how many items are considered in the metric.
+</td>
+</tr><tr>
+<td>
+`**kwargs`
+</td>
+<td>
+Keyword arguments for the metric object.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
+<tr class="alt">
+<td colspan="2">
+A tf.keras.metrics.Metric. See `_RankingMetric` signature for more details.
+</td>
+</tr>
+
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Raises</h2></th></tr>
+
+<tr>
+<td>
+`ValueError`
+</td>
+<td>
+If key is unsupported.
+</td>
+</tr>
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model.md
index bbcc22d..d6d7150 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model.md
@@ -1,3 +1,5 @@
+description: Ranking Model utilities and classes in Keras.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.model" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/create_keras_model.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/create_keras_model.md
index 964ed4c..b66f911 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/create_keras_model.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/model/create_keras_model.md
@@ -1,3 +1,5 @@
+description: Creates a Functional Keras ranking model.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.model.create_keras_model" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/model.py#L9-L56">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network.md
index f42a9d7..797808a 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network.md
@@ -1,3 +1,5 @@
+description: Ranking Networks in Keras.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.network" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network/RankingNetwork.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network/RankingNetwork.md
index 31342b1..764b02c 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network/RankingNetwork.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network/RankingNetwork.md
@@ -1,3 +1,5 @@
+description: Base class for ranking networks in Keras.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.network.RankingNetwork" />
 <meta itemprop="path" content="Stable" />
@@ -23,10 +25,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L85-L216">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -38,7 +39,7 @@ Base class for ranking networks in Keras.
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.network.RankingNetwork(
     context_feature_columns=None, example_feature_columns=None,
-    name='ranking_network', **kwargs
+    name=&#x27;ranking_network&#x27;, **kwargs
 )
 </code></pre>
 
@@ -86,13 +87,32 @@ keyword arguments.
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `context_feature_columns` </td>
-<td>
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`context_feature_columns` </td> <td>
+
+</td> </tr><tr> <td> `dtype` </td> <td> The dtype of the layer weights.
 
-</td> </tr><tr> <td> `dtype` </td> <td> Dtype used by the weights of the layer,
-set in the constructor. </td> </tr><tr> <td> `dynamic` </td> <td> Whether the
-layer is dynamic (eager-only); set in the constructor. </td> </tr><tr> <td>
-`example_feature_columns` </td> <td>
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `example_feature_columns` </td> <td>
 
 </td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
 layer.
@@ -219,6 +239,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -290,9 +317,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -331,11 +359,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -414,9 +442,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -432,7 +461,7 @@ Instance of `TensorShape`, or list of instances of
 
 <h3 id="compute_logits"><code>compute_logits</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L143-L165">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -507,9 +536,10 @@ all entries are valid.
 Computes an output mask tensor.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -556,9 +586,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -623,7 +654,7 @@ if the layer isn't yet built
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L198-L216">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -672,7 +703,7 @@ A RankingNetwork layer.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L188-L196">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -688,6 +719,10 @@ its trained weights) from this configuration.
 The config of a layer does not include connectivity information, nor the layer
 class name. These are handled by `Network` (one layer of abstraction above).
 
+Note that `get_config()` does not guarantee to return a fresh copy of dict every
+time it is called. The callers should make a copy of the returned dict if they
+want to modify it.
+
 <!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
@@ -791,9 +826,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -827,7 +863,7 @@ layer's specifications.
 
 <h3 id="transform"><code>transform</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L118-L141">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -966,9 +1002,10 @@ The original method wrapped such that it enters the module's name scope.
 Wraps `call`, applying pre- and post-processing steps.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -1009,6 +1046,7 @@ Output tensor(s).
     do), its default value will be set to the mask generated for `inputs` by the
     previous layer (if `input` did come from a layer that generated a
     corresponding mask, i.e. if it came from a Keras layer with masking support.
+-   If the layer is not built, the method will call `build`.
 
 <!-- Tabular view -->
  <table class="responsive fixed orange">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network/UnivariateRankingNetwork.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network/UnivariateRankingNetwork.md
index 3907e1c..07d3fee 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network/UnivariateRankingNetwork.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network/UnivariateRankingNetwork.md
@@ -1,3 +1,5 @@
+description: Base class for univariate ranking network.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.network.UnivariateRankingNetwork" />
 <meta itemprop="path" content="Stable" />
@@ -24,10 +26,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L219-L289">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -41,7 +42,7 @@ Inherits From: [`RankingNetwork`](../../../tfr/keras/network/RankingNetwork.md)
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.keras.network.UnivariateRankingNetwork(
     context_feature_columns=None, example_feature_columns=None,
-    name='univariate_ranking_network', **kwargs
+    name=&#x27;univariate_ranking_network&#x27;, **kwargs
 )
 </code></pre>
 
@@ -89,13 +90,32 @@ keyword arguments.
 <tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
 
 <tr> <td> `activity_regularizer` </td> <td> Optional regularizer function for
-the output of this layer. </td> </tr><tr> <td> `context_feature_columns` </td>
-<td>
+the output of this layer. </td> </tr><tr> <td> `compute_dtype` </td> <td> The
+dtype of the layer's computations.
+
+This is equivalent to `Layer.dtype_policy.compute_dtype`. Unless mixed precision
+is used, this is the same as `Layer.dtype`, the dtype of the weights.
+
+Layers automatically cast their inputs to the compute dtype, which causes
+computations and the output to be in the compute dtype as well. This is done by
+the base Layer class in `Layer.__call__`, so you do not have to insert these
+casts if implementing your own layer.
+
+Layers often perform certain internal computations in higher precision when
+`compute_dtype` is float16 or bfloat16 for numeric stability. The output will
+still typically be float16 or bfloat16 in such cases. </td> </tr><tr> <td>
+`context_feature_columns` </td> <td>
+
+</td> </tr><tr> <td> `dtype` </td> <td> The dtype of the layer weights.
 
-</td> </tr><tr> <td> `dtype` </td> <td> Dtype used by the weights of the layer,
-set in the constructor. </td> </tr><tr> <td> `dynamic` </td> <td> Whether the
-layer is dynamic (eager-only); set in the constructor. </td> </tr><tr> <td>
-`example_feature_columns` </td> <td>
+This is equivalent to `Layer.dtype_policy.variable_dtype`. Unless mixed
+precision is used, this is the same as `Layer.compute_dtype`, the dtype of the
+layer's computations. </td> </tr><tr> <td> `dtype_policy` </td> <td> The dtype
+policy associated with this layer.
+
+This is an instance of a `tf.keras.mixed_precision.Policy`. </td> </tr><tr> <td>
+`dynamic` </td> <td> Whether the layer is dynamic (eager-only); set in the
+constructor. </td> </tr><tr> <td> `example_feature_columns` </td> <td>
 
 </td> </tr><tr> <td> `input` </td> <td> Retrieves the input tensor(s) of a
 layer.
@@ -222,6 +242,13 @@ Trainable weights are updated via gradient descent during training.
 </td>
 </tr><tr>
 <td>
+`variable_dtype`
+</td>
+<td>
+Alias of `Layer.dtype`, the dtype of the weights.
+</td>
+</tr><tr>
+<td>
 `weights`
 </td>
 <td>
@@ -293,9 +320,10 @@ model.add_loss(lambda: tf.reduce_mean(d.kernel))
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -334,11 +362,11 @@ model.
 class MyMetricLayer(tf.keras.layers.Layer):
   def __init__(self):
     super(MyMetricLayer, self).__init__(name='my_metric_layer')
-    self.mean = metrics_module.Mean(name='metric_1')
+    self.mean = tf.keras.metrics.Mean(name='metric_1')
 
   def call(self, inputs):
     self.add_metric(self.mean(x))
-    self.add_metric(math_ops.reduce_sum(x), name='metric_2')
+    self.add_metric(tf.reduce_sum(x), name='metric_2')
     return inputs
 ```
 
@@ -417,9 +445,10 @@ layer call.
 This is typically used to create the weights of `Layer` subclasses.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -435,7 +464,7 @@ Instance of `TensorShape`, or list of instances of
 
 <h3 id="compute_logits"><code>compute_logits</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L252-L289">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -527,7 +556,7 @@ Computes an output mask tensor.
 
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -574,9 +603,10 @@ This assumes that the layer will later be used with inputs that match the input
 shape provided here.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -641,7 +671,7 @@ if the layer isn't yet built
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L198-L216">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -690,7 +720,7 @@ A RankingNetwork layer.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L188-L196">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -706,6 +736,10 @@ its trained weights) from this configuration.
 The config of a layer does not include connectivity information, nor the layer
 class name. These are handled by `Network` (one layer of abstraction above).
 
+Note that `get_config()` does not guarantee to return a fresh copy of dict every
+time it is called. The callers should make a copy of the returned dict if they
+want to modify it.
+
 <!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
@@ -770,7 +804,7 @@ Weights values as a list of numpy arrays.
 
 <h3 id="score"><code>score</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L235-L250">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -866,9 +900,10 @@ the bias value. These can be used to set the weights of another Dense layer:
 ```
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -902,7 +937,7 @@ layer's specifications.
 
 <h3 id="transform"><code>transform</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">View
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L118-L141">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -1041,9 +1076,10 @@ The original method wrapped such that it enters the module's name scope.
 Wraps `call`, applying pre- and post-processing steps.
 
 <!-- Tabular view -->
+
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
-<tr><th colspan="2">Arguments</th></tr>
+<tr><th colspan="2">Args</th></tr>
 
 <tr>
 <td>
@@ -1084,6 +1120,7 @@ Output tensor(s).
     do), its default value will be set to the mask generated for `inputs` by the
     previous layer (if `input` did come from a layer that generated a
     corresponding mask, i.e. if it came from a Keras layer with masking support.
+-   If the layer is not built, the method will call `build`.
 
 <!-- Tabular view -->
  <table class="responsive fixed orange">
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network/listwise_scoring.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network/listwise_scoring.md
index 8f4074f..ec10192 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network/listwise_scoring.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/network/listwise_scoring.md
@@ -1,3 +1,5 @@
+description: Listwise scoring op for context and example features.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.keras.network.listwise_scoring" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/network.py#L12-L82">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -83,7 +84,6 @@ is True for a valid example and False for invalid one.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/saved_model.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/saved_model.md
new file mode 100644
index 0000000..b48fbd4
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/saved_model.md
@@ -0,0 +1,26 @@
+description: SavedModel utilities for TF-Ranking.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.saved_model" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# Module: tfr.keras.saved_model
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/saved_model.py">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+SavedModel utilities for TF-Ranking.
+
+## Classes
+
+[`class Signatures`](../../tfr/keras/saved_model/Signatures.md): Defines
+signatures to support regress and predict serving.
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/saved_model/Signatures.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/saved_model/Signatures.md
new file mode 100644
index 0000000..9f31b18
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/saved_model/Signatures.md
@@ -0,0 +1,258 @@
+description: Defines signatures to support regress and predict serving.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.saved_model.Signatures" />
+<meta itemprop="path" content="Stable" />
+<meta itemprop="property" content="__call__"/>
+<meta itemprop="property" content="__init__"/>
+<meta itemprop="property" content="predict_tf_function"/>
+<meta itemprop="property" content="regress_tf_function"/>
+<meta itemprop="property" content="with_name_scope"/>
+</div>
+
+# tfr.keras.saved_model.Signatures
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/saved_model.py#L35-L130">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Defines signatures to support regress and predict serving.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.saved_model.Signatures(
+    model: tf.keras.Model,
+    context_feature_spec: Dict[str, Union[tf.io.FixedLenFeature, tf.io.RaggedFeature]],
+    example_feature_spec: Dict[str, Union[tf.io.FixedLenFeature, tf.io.RaggedFeature]],
+    mask_feature_name: str
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`model`
+</td>
+<td>
+A keras ranking model.
+</td>
+</tr><tr>
+<td>
+`context_feature_spec`
+</td>
+<td>
+(dict) A mapping from feature keys to
+`FixedLenFeature` or `RaggedFeature` values for context in
+`tensorflow.serving.ExampleListWithContext` proto.
+</td>
+</tr><tr>
+<td>
+`example_feature_spec`
+</td>
+<td>
+(dict) A mapping from feature keys to
+`FixedLenFeature` or `Ragged` values for examples in
+`tensorflow.serving.ExampleListWithContext` proto.
+</td>
+</tr><tr>
+<td>
+`mask_feature_name`
+</td>
+<td>
+(str) Name of feature for example list masks.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>
+
+<tr> <td> `name` </td> <td> Returns the name of this module as passed or
+determined in the ctor.
+
+NOTE: This is not the same as the `self.name_scope.name` which includes parent
+module names. </td> </tr><tr> <td> `name_scope` </td> <td> Returns a
+`tf.name_scope` instance for this class. </td> </tr><tr> <td>
+`non_trainable_variables` </td> <td> Sequence of non-trainable variables owned
+by this module and its submodules.
+
+Note: this method uses reflection to find variables on the current instance and
+submodules. For performance reasons you may wish to cache the result of calling
+this method if you don't expect the return value to change. </td> </tr><tr> <td>
+`submodules` </td> <td> Sequence of all sub-modules.
+
+Submodules are modules which are properties of this module, or found as
+properties of modules which are properties of this module (and so on).
+
+```
+>>> a = tf.Module()
+>>> b = tf.Module()
+>>> c = tf.Module()
+>>> a.b = b
+>>> b.c = c
+>>> list(a.submodules) == [b, c]
+True
+>>> list(b.submodules) == [c]
+True
+>>> list(c.submodules) == []
+True
+```
+
+</td> </tr><tr> <td> `trainable_variables` </td> <td> Sequence of trainable
+variables owned by this module and its submodules.
+
+Note: this method uses reflection to find variables on the current instance and
+submodules. For performance reasons you may wish to cache the result of calling
+this method if you don't expect the return value to change. </td> </tr><tr> <td>
+`variables` </td> <td> Sequence of variables owned by this module and its
+submodules.
+
+Note: this method uses reflection to find variables on the current instance
+and submodules. For performance reasons you may wish to cache the result
+of calling this method if you don't expect the return value to change.
+</td>
+</tr>
+</table>
+
+## Methods
+
+<h3 id="predict_tf_function"><code>predict_tf_function</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/saved_model.py#L62-L79">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>predict_tf_function() -> Callable[[tf.Tensor], Dict[str, tf.Tensor]]
+</code></pre>
+
+Makes a tensorflow function for `predict`.
+
+<h3 id="regress_tf_function"><code>regress_tf_function</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/saved_model.py#L81-L100">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>regress_tf_function() -> Callable[[tf.Tensor], Dict[str, tf.Tensor]]
+</code></pre>
+
+Makes a tensorflow function for `regress`.
+
+<h3 id="with_name_scope"><code>with_name_scope</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@classmethod</code>
+<code>with_name_scope(
+    method
+)
+</code></pre>
+
+Decorator to automatically enter the module name scope.
+
+```
+>>> class MyModule(tf.Module):
+...   @tf.Module.with_name_scope
+...   def __call__(self, x):
+...     if not hasattr(self, 'w'):
+...       self.w = tf.Variable(tf.random.normal([x.shape[1], 3]))
+...     return tf.matmul(x, self.w)
+```
+
+Using the above module would produce `tf.Variable`s and `tf.Tensor`s whose names
+included the module name:
+
+```
+>>> mod = MyModule()
+>>> mod(tf.ones([1, 2]))
+<tf.Tensor: shape=(1, 3), dtype=float32, numpy=..., dtype=float32)>
+>>> mod.w
+<tf.Variable 'my_module/Variable:0' shape=(2, 3) dtype=float32,
+numpy=..., dtype=float32)>
+```
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`method`
+</td>
+<td>
+The method to wrap.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+The original method wrapped such that it enters the module's name scope.
+</td>
+</tr>
+
+</table>
+
+<h3 id="__call__"><code>__call__</code></h3>
+
+<a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/saved_model.py#L102-L130">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>__call__(
+    serving_default: str = &#x27;regress&#x27;
+) -> Dict[str, Callable[[tf.Tensor], Dict[str, tf.Tensor]]]
+</code></pre>
+
+Returns a dict of signatures.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`serving_default`
+</td>
+<td>
+Specifies "regress" or "predict" as the serving_default
+signature.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Returns</th></tr>
+<tr class="alt">
+<td colspan="2">
+A dict of signatures.
+</td>
+</tr>
+
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils.md
new file mode 100644
index 0000000..ce8468c
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils.md
@@ -0,0 +1,33 @@
+description: Utils for tfr.keras.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.utils" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# Module: tfr.keras.utils
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/utils.py">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Utils for tfr.keras.
+
+## Functions
+
+[`identity(...)`](../../tfr/keras/utils/identity.md)
+
+[`inverse(...)`](../../tfr/keras/utils/inverse.md)
+
+[`is_greater_equal_1(...)`](../../tfr/keras/utils/is_greater_equal_1.md)
+
+[`log2_inverse(...)`](../../tfr/keras/utils/log2_inverse.md)
+
+[`pow_minus_1(...)`](../../tfr/keras/utils/pow_minus_1.md)
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/identity.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/identity.md
new file mode 100644
index 0000000..cd3267a
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/identity.md
@@ -0,0 +1,25 @@
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.utils.identity" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# tfr.keras.utils.identity
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/utils.py#L10-L12">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.utils.identity(
+    label
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/inverse.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/inverse.md
new file mode 100644
index 0000000..2ca49e4
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/inverse.md
@@ -0,0 +1,25 @@
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.utils.inverse" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# tfr.keras.utils.inverse
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/utils.py#L15-L17">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.utils.inverse(
+    rank
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/is_greater_equal_1.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/is_greater_equal_1.md
new file mode 100644
index 0000000..5e16d1d
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/is_greater_equal_1.md
@@ -0,0 +1,25 @@
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.utils.is_greater_equal_1" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# tfr.keras.utils.is_greater_equal_1
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/utils.py#L30-L32">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.utils.is_greater_equal_1(
+    label
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/log2_inverse.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/log2_inverse.md
new file mode 100644
index 0000000..0e4d7c3
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/log2_inverse.md
@@ -0,0 +1,25 @@
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.utils.log2_inverse" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# tfr.keras.utils.log2_inverse
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/utils.py#L25-L27">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.utils.log2_inverse(
+    rank
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/pow_minus_1.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/pow_minus_1.md
new file mode 100644
index 0000000..c65190c
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/keras/utils/pow_minus_1.md
@@ -0,0 +1,25 @@
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.utils.pow_minus_1" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# tfr.keras.utils.pow_minus_1
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/utils.py#L20-L22">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.utils.pow_minus_1(
+    label
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/losses.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/losses.md
index 6a25975..504781d 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/losses.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/losses.md
@@ -1,3 +1,5 @@
+description: Defines ranking losses as TF ops.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.losses" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/RankingLossKey.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/RankingLossKey.md
index 6614b99..839b960 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/RankingLossKey.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/RankingLossKey.md
@@ -1,3 +1,5 @@
+description: Ranking loss key strings.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.losses.RankingLossKey" />
 <meta itemprop="path" content="Stable" />
@@ -20,10 +22,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses.py#L17-L32">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -34,27 +35,102 @@ Ranking loss key strings.
 
 <!-- Placeholder for "Used in" -->
 
-## Class Variables
-
-*   `APPROX_MRR_LOSS = 'approx_mrr_loss'` <a id="APPROX_MRR_LOSS"></a>
-*   `APPROX_NDCG_LOSS = 'approx_ndcg_loss'` <a id="APPROX_NDCG_LOSS"></a>
-*   `GUMBEL_APPROX_NDCG_LOSS = 'gumbel_approx_ndcg_loss'`
-    <a id="GUMBEL_APPROX_NDCG_LOSS"></a>
-*   `GUMBEL_NEURAL_SORT_CROSS_ENTROPY_LOSS =
-    'gumbel_neural_sort_cross_entropy_loss'`
-    <a id="GUMBEL_NEURAL_SORT_CROSS_ENTROPY_LOSS"></a>
-*   `LIST_MLE_LOSS = 'list_mle_loss'` <a id="LIST_MLE_LOSS"></a>
-*   `MEAN_SQUARED_LOSS = 'mean_squared_loss'` <a id="MEAN_SQUARED_LOSS"></a>
-*   `NEURAL_SORT_CROSS_ENTROPY_LOSS = 'neural_sort_cross_entropy_loss'`
-    <a id="NEURAL_SORT_CROSS_ENTROPY_LOSS"></a>
-*   `PAIRWISE_HINGE_LOSS = 'pairwise_hinge_loss'`
-    <a id="PAIRWISE_HINGE_LOSS"></a>
-*   `PAIRWISE_LOGISTIC_LOSS = 'pairwise_logistic_loss'`
-    <a id="PAIRWISE_LOGISTIC_LOSS"></a>
-*   `PAIRWISE_SOFT_ZERO_ONE_LOSS = 'pairwise_soft_zero_one_loss'`
-    <a id="PAIRWISE_SOFT_ZERO_ONE_LOSS"></a>
-*   `SIGMOID_CROSS_ENTROPY_LOSS = 'sigmoid_cross_entropy_loss'`
-    <a id="SIGMOID_CROSS_ENTROPY_LOSS"></a>
-*   `SOFTMAX_LOSS = 'softmax_loss'` <a id="SOFTMAX_LOSS"></a>
-*   `UNIQUE_SOFTMAX_LOSS = 'unique_softmax_loss'`
-    <a id="UNIQUE_SOFTMAX_LOSS"></a>
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Class Variables</h2></th></tr>
+
+<tr>
+<td>
+APPROX_MRR_LOSS<a id="APPROX_MRR_LOSS"></a>
+</td>
+<td>
+`'approx_mrr_loss'`
+</td>
+</tr><tr>
+<td>
+APPROX_NDCG_LOSS<a id="APPROX_NDCG_LOSS"></a>
+</td>
+<td>
+`'approx_ndcg_loss'`
+</td>
+</tr><tr>
+<td>
+GUMBEL_APPROX_NDCG_LOSS<a id="GUMBEL_APPROX_NDCG_LOSS"></a>
+</td>
+<td>
+`'gumbel_approx_ndcg_loss'`
+</td>
+</tr><tr>
+<td>
+GUMBEL_NEURAL_SORT_CROSS_ENTROPY_LOSS<a id="GUMBEL_NEURAL_SORT_CROSS_ENTROPY_LOSS"></a>
+</td>
+<td>
+`'gumbel_neural_sort_cross_entropy_loss'`
+</td>
+</tr><tr>
+<td>
+LIST_MLE_LOSS<a id="LIST_MLE_LOSS"></a>
+</td>
+<td>
+`'list_mle_loss'`
+</td>
+</tr><tr>
+<td>
+MEAN_SQUARED_LOSS<a id="MEAN_SQUARED_LOSS"></a>
+</td>
+<td>
+`'mean_squared_loss'`
+</td>
+</tr><tr>
+<td>
+NEURAL_SORT_CROSS_ENTROPY_LOSS<a id="NEURAL_SORT_CROSS_ENTROPY_LOSS"></a>
+</td>
+<td>
+`'neural_sort_cross_entropy_loss'`
+</td>
+</tr><tr>
+<td>
+PAIRWISE_HINGE_LOSS<a id="PAIRWISE_HINGE_LOSS"></a>
+</td>
+<td>
+`'pairwise_hinge_loss'`
+</td>
+</tr><tr>
+<td>
+PAIRWISE_LOGISTIC_LOSS<a id="PAIRWISE_LOGISTIC_LOSS"></a>
+</td>
+<td>
+`'pairwise_logistic_loss'`
+</td>
+</tr><tr>
+<td>
+PAIRWISE_SOFT_ZERO_ONE_LOSS<a id="PAIRWISE_SOFT_ZERO_ONE_LOSS"></a>
+</td>
+<td>
+`'pairwise_soft_zero_one_loss'`
+</td>
+</tr><tr>
+<td>
+SIGMOID_CROSS_ENTROPY_LOSS<a id="SIGMOID_CROSS_ENTROPY_LOSS"></a>
+</td>
+<td>
+`'sigmoid_cross_entropy_loss'`
+</td>
+</tr><tr>
+<td>
+SOFTMAX_LOSS<a id="SOFTMAX_LOSS"></a>
+</td>
+<td>
+`'softmax_loss'`
+</td>
+</tr><tr>
+<td>
+UNIQUE_SOFTMAX_LOSS<a id="UNIQUE_SOFTMAX_LOSS"></a>
+</td>
+<td>
+`'unique_softmax_loss'`
+</td>
+</tr>
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/create_ndcg_lambda_weight.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/create_ndcg_lambda_weight.md
index de23c15..a50c3c3 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/create_ndcg_lambda_weight.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/create_ndcg_lambda_weight.md
@@ -1,3 +1,5 @@
+description: Creates _LambdaWeight for NDCG metric.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.losses.create_ndcg_lambda_weight" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses.py#L256-L263">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/create_p_list_mle_lambda_weight.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/create_p_list_mle_lambda_weight.md
index 685f771..8b8035c 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/create_p_list_mle_lambda_weight.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/create_p_list_mle_lambda_weight.md
@@ -1,3 +1,5 @@
+description: Creates _LambdaWeight based on Position-Aware ListMLE paper.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.losses.create_p_list_mle_lambda_weight" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses.py#L276-L290">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/create_reciprocal_rank_lambda_weight.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/create_reciprocal_rank_lambda_weight.md
index 207d71f..1491998 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/create_reciprocal_rank_lambda_weight.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/create_reciprocal_rank_lambda_weight.md
@@ -1,3 +1,5 @@
+description: Creates _LambdaWeight for MRR-like metric.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.losses.create_reciprocal_rank_lambda_weight" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses.py#L266-L273">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/make_loss_fn.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/make_loss_fn.md
index 8a8f18f..265b194 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/make_loss_fn.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/make_loss_fn.md
@@ -1,3 +1,5 @@
+description: Makes a loss function using a single loss or multiple losses.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.losses.make_loss_fn" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses.py#L35-L181">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/make_loss_metric_fn.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/make_loss_metric_fn.md
index eec9917..8e82687 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/make_loss_metric_fn.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/losses/make_loss_metric_fn.md
@@ -1,3 +1,5 @@
+description: Factory method to create a metric based on a loss.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.losses.make_loss_metric_fn" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses.py#L184-L253">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics.md
index 8491d4a..2f7891f 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics.md
@@ -1,3 +1,5 @@
+description: Defines ranking metrics as TF ops.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.metrics" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
@@ -37,6 +38,9 @@ Computes alpha discounted cumulative gain (alpha-DCG).
 [`average_relevance_position(...)`](../tfr/metrics/average_relevance_position.md):
 Computes average relevance position (ARP).
 
+[`binary_preference(...)`](../tfr/metrics/binary_preference.md): Computes binary
+preference (BPref).
+
 [`compute_mean(...)`](../tfr/metrics/compute_mean.md): Returns the mean of the
 specified metric given the inputs.
 
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/RankingMetricKey.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/RankingMetricKey.md
index 6eb5b07..b06c0b8 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/RankingMetricKey.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/RankingMetricKey.md
@@ -1,8 +1,11 @@
+description: Ranking metric key strings.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.metrics.RankingMetricKey" />
 <meta itemprop="path" content="Stable" />
 <meta itemprop="property" content="ALPHA_DCG"/>
 <meta itemprop="property" content="ARP"/>
+<meta itemprop="property" content="BPREF"/>
 <meta itemprop="property" content="DCG"/>
 <meta itemprop="property" content="MAP"/>
 <meta itemprop="property" content="MRR"/>
@@ -16,10 +19,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py#L23-L53">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -30,15 +32,81 @@ Ranking metric key strings.
 
 <!-- Placeholder for "Used in" -->
 
-## Class Variables
-
-*   `ALPHA_DCG = 'alpha_dcg'` <a id="ALPHA_DCG"></a>
-*   `ARP = 'arp'` <a id="ARP"></a>
-*   `DCG = 'dcg'` <a id="DCG"></a>
-*   `MAP = 'map'` <a id="MAP"></a>
-*   `MRR = 'mrr'` <a id="MRR"></a>
-*   `NDCG = 'ndcg'` <a id="NDCG"></a>
-*   `ORDERED_PAIR_ACCURACY = 'ordered_pair_accuracy'`
-    <a id="ORDERED_PAIR_ACCURACY"></a>
-*   `PRECISION = 'precision'` <a id="PRECISION"></a>
-*   `PRECISION_IA = 'precision_ia'` <a id="PRECISION_IA"></a>
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Class Variables</h2></th></tr>
+
+<tr>
+<td>
+ALPHA_DCG<a id="ALPHA_DCG"></a>
+</td>
+<td>
+`'alpha_dcg'`
+</td>
+</tr><tr>
+<td>
+ARP<a id="ARP"></a>
+</td>
+<td>
+`'arp'`
+</td>
+</tr><tr>
+<td>
+BPREF<a id="BPREF"></a>
+</td>
+<td>
+`'bpref'`
+</td>
+</tr><tr>
+<td>
+DCG<a id="DCG"></a>
+</td>
+<td>
+`'dcg'`
+</td>
+</tr><tr>
+<td>
+MAP<a id="MAP"></a>
+</td>
+<td>
+`'map'`
+</td>
+</tr><tr>
+<td>
+MRR<a id="MRR"></a>
+</td>
+<td>
+`'mrr'`
+</td>
+</tr><tr>
+<td>
+NDCG<a id="NDCG"></a>
+</td>
+<td>
+`'ndcg'`
+</td>
+</tr><tr>
+<td>
+ORDERED_PAIR_ACCURACY<a id="ORDERED_PAIR_ACCURACY"></a>
+</td>
+<td>
+`'ordered_pair_accuracy'`
+</td>
+</tr><tr>
+<td>
+PRECISION<a id="PRECISION"></a>
+</td>
+<td>
+`'precision'`
+</td>
+</tr><tr>
+<td>
+PRECISION_IA<a id="PRECISION_IA"></a>
+</td>
+<td>
+`'precision_ia'`
+</td>
+</tr>
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/alpha_discounted_cumulative_gain.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/alpha_discounted_cumulative_gain.md
index 7d5b370..a68dc00 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/alpha_discounted_cumulative_gain.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/alpha_discounted_cumulative_gain.md
@@ -1,3 +1,5 @@
+description: Computes alpha discounted cumulative gain (alpha-DCG).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.metrics.alpha_discounted_cumulative_gain" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py#L457-L500">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/average_relevance_position.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/average_relevance_position.md
index 62f2875..7df507b 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/average_relevance_position.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/average_relevance_position.md
@@ -1,3 +1,5 @@
+description: Computes average relevance position (ARP).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.metrics.average_relevance_position" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py#L277-L301">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/binary_preference.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/binary_preference.md
new file mode 100644
index 0000000..213e221
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/binary_preference.md
@@ -0,0 +1,102 @@
+description: Computes binary preference (BPref).
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.metrics.binary_preference" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# tfr.metrics.binary_preference
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py#L531-L567">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Computes binary preference (BPref).
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.metrics.binary_preference(
+    labels, predictions, weights=None, topn=None, name=None, use_trec_version=True
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+
+The implementation of BPref is based on the desciption in the following:
+https://trec.nist.gov/pubs/trec15/appendices/CE.MEASURES06.pdf BPref = 1 / R
+SUM_r(1 - |n ranked higher than r| / min(R, N))
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`labels`
+</td>
+<td>
+A `Tensor` of the same shape as `predictions`. A value >= 1 means a
+relevant example.
+</td>
+</tr><tr>
+<td>
+`predictions`
+</td>
+<td>
+A `Tensor` with shape [batch_size, list_size]. Each value is
+the ranking score of the corresponding example.
+</td>
+</tr><tr>
+<td>
+`weights`
+</td>
+<td>
+A `Tensor` of the same shape of predictions or [batch_size, 1]. The
+former case is per-example and the latter case is per-list.
+</td>
+</tr><tr>
+<td>
+`topn`
+</td>
+<td>
+A cutoff for how many examples to consider for this metric.
+</td>
+</tr><tr>
+<td>
+`name`
+</td>
+<td>
+A string used as the name for this metric.
+</td>
+</tr><tr>
+<td>
+`use_trec_version`
+</td>
+<td>
+A boolean to choose the version of the formula to use.
+If False, than the alternative BPref formula will be used:
+BPref = 1 / R SUM_r(1 - |n ranked higher than r| / R)
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
+<tr class="alt">
+<td colspan="2">
+A metric for binary preference metric of the batch.
+</td>
+</tr>
+
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/compute_mean.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/compute_mean.md
index 350bd6a..d5cb14e 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/compute_mean.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/compute_mean.md
@@ -1,3 +1,5 @@
+description: Returns the mean of the specified metric given the inputs.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.metrics.compute_mean" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py#L56-L95">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/discounted_cumulative_gain.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/discounted_cumulative_gain.md
index de927f0..151a525 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/discounted_cumulative_gain.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/discounted_cumulative_gain.md
@@ -1,3 +1,5 @@
+description: Computes discounted cumulative gain (DCG).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.metrics.discounted_cumulative_gain" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py#L426-L454">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/eval_metric.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/eval_metric.md
index d4d7fc3..70b7ac9 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/eval_metric.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/eval_metric.md
@@ -1,3 +1,5 @@
+description: A stand-alone method to evaluate metrics on ranked results.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.metrics.eval_metric" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py#L570-L621">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/make_ranking_metric_fn.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/make_ranking_metric_fn.md
index 372cfbc..c306413 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/make_ranking_metric_fn.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/make_ranking_metric_fn.md
@@ -1,3 +1,5 @@
+description: Factory method to create a ranking metric function.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.metrics.make_ranking_metric_fn" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py#L98-L245">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/mean_average_precision.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/mean_average_precision.md
index e0ce65b..883eb8f 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/mean_average_precision.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/mean_average_precision.md
@@ -1,3 +1,5 @@
+description: Computes mean average precision (MAP).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.metrics.mean_average_precision" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py#L329-L359">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/mean_reciprocal_rank.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/mean_reciprocal_rank.md
index 41215a8..da312f0 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/mean_reciprocal_rank.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/mean_reciprocal_rank.md
@@ -1,3 +1,5 @@
+description: Computes mean reciprocal rank (MRR).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.metrics.mean_reciprocal_rank" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py#L248-L274">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/normalized_discounted_cumulative_gain.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/normalized_discounted_cumulative_gain.md
index b492d54..49721d3 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/normalized_discounted_cumulative_gain.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/normalized_discounted_cumulative_gain.md
@@ -1,3 +1,5 @@
+description: Computes normalized discounted cumulative gain (NDCG).
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.metrics.normalized_discounted_cumulative_gain" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py#L387-L423">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/ordered_pair_accuracy.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/ordered_pair_accuracy.md
index dc2096a..522bcdd 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/ordered_pair_accuracy.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/ordered_pair_accuracy.md
@@ -1,3 +1,5 @@
+description: Computes the percentage of correctly ordered pair.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.metrics.ordered_pair_accuracy" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py#L503-L528">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/precision.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/precision.md
index 2091549..68d720c 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/precision.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/precision.md
@@ -1,3 +1,5 @@
+description: Computes precision as weighted average of relevant examples.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.metrics.precision" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py#L304-L326">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/precision_ia.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/precision_ia.md
index cfdca03..fd5f12a 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/precision_ia.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/metrics/precision_ia.md
@@ -1,3 +1,6 @@
+description: Computes Intent-Aware Precision as weighted average of relevant
+examples.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.metrics.precision_ia" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +10,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/metrics.py#L362-L384">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/model.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/model.md
index e3488f6..09dd4ee 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/model.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/model.md
@@ -1,3 +1,5 @@
+description: Tensorflow Ranking model library.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.model" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/model.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/model/make_groupwise_ranking_fn.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/model/make_groupwise_ranking_fn.md
index 5c345c2..c4d1ae7 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/model/make_groupwise_ranking_fn.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/model/make_groupwise_ranking_fn.md
@@ -1,3 +1,6 @@
+description: Builds an Estimator model_fn for groupwise comparison ranking
+models.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.model.make_groupwise_ranking_fn" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +10,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/model.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/model.py#L439-L458">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils.md
index 853e202..8ecbf69 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils.md
@@ -1,3 +1,5 @@
+description: Utility functions for ranking library.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.utils" />
 <meta itemprop="path" content="Stable" />
@@ -7,8 +9,7 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
   <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
@@ -24,6 +25,9 @@ Utility functions for ranking library.
 [`de_noise(...)`](../tfr/utils/de_noise.md): Returns a float `Tensor` as the
 de-noised `counts`.
 
+[`gather_per_row(...)`](../tfr/utils/gather_per_row.md): Gathers the values from
+input tensor based on per-row indices.
+
 [`is_label_valid(...)`](../tfr/utils/is_label_valid.md): Returns a boolean
 `Tensor` for label validity.
 
@@ -33,6 +37,9 @@ Organizes indices in such a way that valid items appear first.
 [`padded_nd_indices(...)`](../tfr/utils/padded_nd_indices.md): Pads the invalid
 entries by valid ones and returns the nd_indices.
 
+[`ragged_to_dense(...)`](../tfr/utils/ragged_to_dense.md): Converts given inputs
+from ragged tensors to dense tensors.
+
 [`reshape_first_ndims(...)`](../tfr/utils/reshape_first_ndims.md): Reshapes the
 first n dims of the input `tensor` to `new shape`.
 
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/de_noise.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/de_noise.md
index d68b989..ed3e673 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/de_noise.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/de_noise.md
@@ -1,3 +1,5 @@
+description: Returns a float Tensor as the de-noised counts.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.utils.de_noise" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py#L340-L399">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/gather_per_row.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/gather_per_row.md
new file mode 100644
index 0000000..3163285
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/gather_per_row.md
@@ -0,0 +1,79 @@
+description: Gathers the values from input tensor based on per-row indices.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.utils.gather_per_row" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# tfr.utils.gather_per_row
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py#L33-L56">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Gathers the values from input tensor based on per-row indices.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.utils.gather_per_row(
+    inputs, indices
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+
+#### Example Usage:
+
+```python
+scores = [[1., 3., 2.], [1., 2., 3.]]
+indices = [[1, 2], [2, 1]]
+tfr.utils.gather_per_row(scores, indices)
+```
+
+Returns [[3., 2.], [3., 2.]]
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`inputs`
+</td>
+<td>
+(tf.Tensor) A tensor of shape [batch_size, list_size] or
+[batch_size, list_size, feature_dims].
+</td>
+</tr><tr>
+<td>
+`indices`
+</td>
+<td>
+(tf.Tensor) A tensor of shape [batch_size, size] of positions to
+gather inputs from. Each index corresponds to a row entry in input_tensor.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
+<tr class="alt">
+<td colspan="2">
+A tensor of values gathered from inputs, of shape [batch_size, size] or
+[batch_size, size, feature_dims], depending on whether the input was 2D or
+3D.
+</td>
+</tr>
+
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/is_label_valid.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/is_label_valid.md
index d9b12c9..d5de15d 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/is_label_valid.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/is_label_valid.md
@@ -1,3 +1,5 @@
+description: Returns a boolean Tensor for label validity.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.utils.is_label_valid" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py#L59-L62">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/organize_valid_indices.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/organize_valid_indices.md
index 96d5727..e08bd08 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/organize_valid_indices.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/organize_valid_indices.md
@@ -1,3 +1,5 @@
+description: Organizes indices in such a way that valid items appear first.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.utils.organize_valid_indices" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py#L184-L216">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/padded_nd_indices.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/padded_nd_indices.md
index 4c19619..2ebf26b 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/padded_nd_indices.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/padded_nd_indices.md
@@ -1,3 +1,5 @@
+description: Pads the invalid entries by valid ones and returns the nd_indices.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.utils.padded_nd_indices" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py#L289-L337">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/ragged_to_dense.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/ragged_to_dense.md
new file mode 100644
index 0000000..9827f1a
--- /dev/null
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/ragged_to_dense.md
@@ -0,0 +1,75 @@
+description: Converts given inputs from ragged tensors to dense tensors.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.utils.ragged_to_dense" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# tfr.utils.ragged_to_dense
+
+<!-- Insert buttons and diff -->
+
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
+<td>
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py#L402-L423">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Converts given inputs from ragged tensors to dense tensors.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.utils.ragged_to_dense(
+    labels, predictions, weights
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>
+
+<tr>
+<td>
+`labels`
+</td>
+<td>
+A `tf.RaggedTensor` of the same shape as `predictions` representing
+relevance.
+</td>
+</tr><tr>
+<td>
+`predictions`
+</td>
+<td>
+A `tf.RaggedTensor` with shape [batch_size, (list_size)]. Each
+value is the ranking score of the corresponding example.
+</td>
+</tr><tr>
+<td>
+`weights`
+</td>
+<td>
+An optional `tf.RaggedTensor` of the same shape of predictions or a
+`tf.Tensor` of shape [batch_size, 1]. The former case is per-example and
+the latter case is per-list.
+</td>
+</tr>
+</table>
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">Returns</h2></th></tr>
+<tr class="alt">
+<td colspan="2">
+A tuple (labels, predictions, weights, mask) of dense `tf.Tensor`s.
+</td>
+</tr>
+
+</table>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/reshape_first_ndims.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/reshape_first_ndims.md
index 1ce56b5..44a87b5 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/reshape_first_ndims.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/reshape_first_ndims.md
@@ -1,3 +1,5 @@
+description: Reshapes the first n dims of the input tensor to new shape.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.utils.reshape_first_ndims" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py#L219-L237">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/reshape_to_2d.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/reshape_to_2d.md
index 5808878..8b0bdc9 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/reshape_to_2d.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/reshape_to_2d.md
@@ -1,3 +1,5 @@
+description: Converts the given tensor to a 2-D Tensor.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.utils.reshape_to_2d" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py#L240-L250">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/shuffle_valid_indices.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/shuffle_valid_indices.md
index 6d661a4..f957294 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/shuffle_valid_indices.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/shuffle_valid_indices.md
@@ -1,3 +1,5 @@
+description: Returns a shuffle of indices with valid ones on top.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.utils.shuffle_valid_indices" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py#L179-L181">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/sort_by_scores.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/sort_by_scores.md
index b9dabb8..9d120ce 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/sort_by_scores.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/sort_by_scores.md
@@ -1,3 +1,5 @@
+description: Sorts list of features according to per-example scores.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.utils.sort_by_scores" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py#L96-L145">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -21,7 +22,7 @@ Sorts list of features according to per-example scores.
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfr.utils.sort_by_scores(
-    scores, features_list, topn=None, shuffle_ties=True, seed=None
+    scores, features_list, topn=None, shuffle_ties=True, seed=None, mask=None
 )
 </code></pre>
 
@@ -70,6 +71,15 @@ A boolean. If True, randomly shuffle before the sorting.
 <td>
 The ops-level random seed used when `shuffle_ties` is True.
 </td>
+</tr><tr>
+<td>
+`mask`
+</td>
+<td>
+An optional `Tensor` of shape [batch_size, list_size] representing
+which entries are valid for sorting. Invalid entries will be pushed to the
+end.
+</td>
 </tr>
 </table>
 
diff --git a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/sorted_ranks.md b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/sorted_ranks.md
index b9c8773..a49c089 100644
--- a/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/sorted_ranks.md
+++ b/tensorflow_ranking/g3doc/api_docs/python/tfr/utils/sorted_ranks.md
@@ -1,3 +1,5 @@
+description: Returns an int Tensor as the ranks (1-based) after sorting scores.
+
 <div itemscope itemtype="http://developers.google.com/ReferenceObject">
 <meta itemprop="name" content="tfr.utils.sorted_ranks" />
 <meta itemprop="path" content="Stable" />
@@ -7,10 +9,9 @@
 
 <!-- Insert buttons and diff -->
 
-<table class="tfo-notebook-buttons tfo-api" align="left">
-
+<table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/utils.py#L148-L176">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/tensorflow_ranking/python/version.py b/tensorflow_ranking/python/version.py
index e2cadd9..8685c4f 100644
--- a/tensorflow_ranking/python/version.py
+++ b/tensorflow_ranking/python/version.py
@@ -17,7 +17,7 @@
 # We follow Semantic Versioning (https://semver.org/)
 _MAJOR_VERSION = '0'
 _MINOR_VERSION = '3'
-_PATCH_VERSION = '2'
+_PATCH_VERSION = '3'
 
 # When building releases, we can update this value on the release branch to
 # reflect the current release candidate ('rc0', 'rc1') or, finally, the official
diff --git a/tensorflow_ranking/tools/pip_package/setup.py b/tensorflow_ranking/tools/pip_package/setup.py
index a4988ae..784947e 100644
--- a/tensorflow_ranking/tools/pip_package/setup.py
+++ b/tensorflow_ranking/tools/pip_package/setup.py
@@ -28,11 +28,11 @@
 # This version string is semver compatible, but incompatible with pip.
 # For pip, we will remove all '-' characters from this string, and use the
 # result for pip.
-_VERSION = '0.3.2'
+_VERSION = '0.3.3'
 
 REQUIRED_PACKAGES = [
     'absl-py >= 0.1.6', 'numpy >= 1.13.3', 'six >= 1.10.0',
-    'tensorflow-serving-api >= 2.0.0, < 3.0.0', 'tf-models-official >= 2.3.0'
+    'tensorflow-serving-api >= 2.0.0, < 3.0.0', 'tf-models-official >= 2.4.0'
 ]
 
 project_name = 'tensorflow_ranking'