From c46cede726fd453e0aaa6097871d23dc8e465bdc Mon Sep 17 00:00:00 2001
From: Le Yan <lyyanle@google.com>
Date: Fri, 2 Jun 2023 19:08:25 -0700
Subject: [PATCH] Update version for release 0.5.3

PiperOrigin-RevId: 537460867
---
 docs/api_docs/python/tfr/_api_cache.json      | 153 ++++++++++---
 docs/api_docs/python/tfr/_toc.yaml            |   8 +
 docs/api_docs/python/tfr/all_symbols.md       |   4 +
 docs/api_docs/python/tfr/api_report.pb        | Bin 70893 -> 74368 bytes
 .../tfr/data/parse_from_example_in_example.md |  19 +-
 .../tfr/data/parse_from_example_list.md       |  19 +-
 .../tfr/data/parse_from_sequence_example.md   |  15 +-
 .../tfr/extension/premade/TFRBertConfig.md    |  18 +-
 .../extension/premade/TFRBertDataConfig.md    |  18 +-
 .../extension/premade/TFRBertModelConfig.md   |   2 -
 .../tfr/extension/premade/TensorDict.md       |   1 -
 .../python/tfr/extension/task/FeatureSpec.md  |   1 -
 .../tfr/extension/task/RankingDataConfig.md   |  18 +-
 .../tfr/extension/task/RankingTaskConfig.md   |  18 +-
 .../python/tfr/keras/layers/Bilinear.md       | 202 +++++++++--------
 .../python/tfr/keras/layers/ConcatFeatures.md | 207 +++++++++---------
 .../layers/DocumentInteractionAttention.md    | 207 +++++++++---------
 .../python/tfr/keras/layers/FlattenList.md    | 207 +++++++++---------
 .../python/tfr/keras/layers/GAMLayer.md       | 207 +++++++++---------
 .../python/tfr/keras/layers/RestoreList.md    | 207 +++++++++---------
 .../tfr/keras/layers/SelfAttentionMask.md     | 207 +++++++++---------
 docs/api_docs/python/tfr/keras/losses.md      |   6 +
 .../python/tfr/keras/losses/ApproxMRRLoss.md  |  23 +-
 .../python/tfr/keras/losses/ApproxNDCGLoss.md |  23 +-
 .../python/tfr/keras/losses/ClickEMLoss.md    |   6 +-
 .../tfr/keras/losses/CoupledRankDistilLoss.md |   9 +-
 .../tfr/keras/losses/DCGLambdaWeight.md       |   4 +-
 .../tfr/keras/losses/GumbelApproxNDCGLoss.md  |  23 +-
 .../tfr/keras/losses/LabelDiffLambdaWeight.md |   4 +-
 .../tfr/keras/losses/ListMLELambdaWeight.md   |   8 +-
 .../python/tfr/keras/losses/ListMLELoss.md    |   8 +-
 .../tfr/keras/losses/MeanSquaredLoss.md       |   6 +-
 .../python/tfr/keras/losses/MixtureEMLoss.md  |   8 +-
 .../tfr/keras/losses/NDCGLambdaWeight.md      |   4 +-
 .../tfr/keras/losses/NDCGLambdaWeightV2.md    |   4 +-
 .../python/tfr/keras/losses/OrdinalLoss.md    |  21 +-
 .../tfr/keras/losses/PairwiseHingeLoss.md     |   8 +-
 .../tfr/keras/losses/PairwiseLogisticLoss.md  |   8 +-
 .../tfr/keras/losses/PairwiseMSELoss.md       |   8 +-
 .../keras/losses/PairwiseSoftZeroOneLoss.md   |   8 +-
 .../tfr/keras/losses/PrecisionLambdaWeight.md |   6 +-
 .../python/tfr/keras/losses/RankingLossKey.md |  12 +-
 .../keras/losses/SigmoidCrossEntropyLoss.md   |  21 +-
 .../python/tfr/keras/losses/SoftmaxLoss.md    |   8 +-
 .../tfr/keras/losses/UniqueSoftmaxLoss.md     |  23 +-
 .../tfr/keras/losses/YetiDCGLambdaWeight.md   | 110 ++++++++++
 .../tfr/keras/losses/YetiLogisticLoss.md      | 204 +++++++++++++++++
 docs/api_docs/python/tfr/keras/losses/get.md  |   2 +-
 .../python/tfr/keras/metrics/ARPMetric.md     | 207 +++++++++---------
 .../tfr/keras/metrics/AlphaDCGMetric.md       | 207 +++++++++---------
 .../python/tfr/keras/metrics/DCGMetric.md     | 207 +++++++++---------
 .../python/tfr/keras/metrics/HitsMetric.md    | 207 +++++++++---------
 .../python/tfr/keras/metrics/MRRMetric.md     | 207 +++++++++---------
 .../metrics/MeanAveragePrecisionMetric.md     | 207 +++++++++---------
 .../python/tfr/keras/metrics/NDCGMetric.md    | 207 +++++++++---------
 .../python/tfr/keras/metrics/OPAMetric.md     | 207 +++++++++---------
 .../tfr/keras/metrics/PrecisionIAMetric.md    | 207 +++++++++---------
 .../tfr/keras/metrics/PrecisionMetric.md      | 207 +++++++++---------
 .../python/tfr/keras/metrics/RecallMetric.md  | 207 +++++++++---------
 .../pipeline/DatasetHparams/dataset_reader.md |  67 ++++--
 docs/api_docs/python/tfr/keras/utils.md       |   4 +
 .../python/tfr/keras/utils/GainFunction.md    |   1 -
 .../keras/utils/deserialize_keras_object.md   |  28 +++
 .../python/tfr/keras/utils/identity.md        |   2 +-
 .../python/tfr/keras/utils/inverse.md         |   2 +-
 .../tfr/keras/utils/is_greater_equal_1.md     |   2 +-
 .../python/tfr/keras/utils/log2_inverse.md    |   2 +-
 .../python/tfr/keras/utils/pow_minus_1.md     |   2 +-
 .../tfr/keras/utils/serialize_keras_object.md |  25 +++
 .../python/tfr/keras/utils/symmetric_log1p.md |   2 +-
 .../api_docs/python/tfr/utils/LossFunction.md |   1 -
 .../tfr/utils/parse_keys_and_weights.md       |   2 -
 72 files changed, 2745 insertions(+), 1985 deletions(-)
 create mode 100644 docs/api_docs/python/tfr/keras/losses/YetiDCGLambdaWeight.md
 create mode 100644 docs/api_docs/python/tfr/keras/losses/YetiLogisticLoss.md
 create mode 100644 docs/api_docs/python/tfr/keras/utils/deserialize_keras_object.md
 create mode 100644 docs/api_docs/python/tfr/keras/utils/serialize_keras_object.md

diff --git a/docs/api_docs/python/tfr/_api_cache.json b/docs/api_docs/python/tfr/_api_cache.json
index 8ed2641..3e48146 100644
--- a/docs/api_docs/python/tfr/_api_cache.json
+++ b/docs/api_docs/python/tfr/_api_cache.json
@@ -300,7 +300,6 @@
     "tfr.keras.layers.ConcatFeatures.__new__": "tfr.keras.layers.Bilinear.__new__",
     "tfr.keras.layers.ConcatFeatures.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.layers.ConcatFeatures.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.layers.ConcatFeatures.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.layers.ConcatFeatures.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.layers.ConcatFeatures.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
     "tfr.keras.layers.ConcatFeatures.compute_mask": "tfr.keras.layers.Bilinear.compute_mask",
@@ -312,12 +311,14 @@
     "tfr.keras.layers.ConcatFeatures.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.layers.ConcatFeatures.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.layers.ConcatFeatures.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.layers.ConcatFeatures.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.layers.ConcatFeatures.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.layers.ConcatFeatures.metrics": "tfr.keras.layers.Bilinear.metrics",
     "tfr.keras.layers.ConcatFeatures.name": "tfr.keras.layers.Bilinear.name",
     "tfr.keras.layers.ConcatFeatures.name_scope": "tfr.keras.layers.Bilinear.name_scope",
     "tfr.keras.layers.ConcatFeatures.non_trainable_weights": "tfr.keras.layers.Bilinear.non_trainable_weights",
     "tfr.keras.layers.ConcatFeatures.output": "tfr.keras.layers.Bilinear.output",
+    "tfr.keras.layers.ConcatFeatures.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.layers.ConcatFeatures.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.layers.ConcatFeatures.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.layers.ConcatFeatures.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -335,7 +336,6 @@
     "tfr.keras.layers.DocumentInteractionAttention.__new__": "tfr.keras.layers.Bilinear.__new__",
     "tfr.keras.layers.DocumentInteractionAttention.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.layers.DocumentInteractionAttention.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.layers.DocumentInteractionAttention.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.layers.DocumentInteractionAttention.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.layers.DocumentInteractionAttention.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
     "tfr.keras.layers.DocumentInteractionAttention.compute_mask": "tfr.keras.layers.Bilinear.compute_mask",
@@ -348,12 +348,14 @@
     "tfr.keras.layers.DocumentInteractionAttention.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.layers.DocumentInteractionAttention.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.layers.DocumentInteractionAttention.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.layers.DocumentInteractionAttention.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.layers.DocumentInteractionAttention.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.layers.DocumentInteractionAttention.metrics": "tfr.keras.layers.Bilinear.metrics",
     "tfr.keras.layers.DocumentInteractionAttention.name": "tfr.keras.layers.Bilinear.name",
     "tfr.keras.layers.DocumentInteractionAttention.name_scope": "tfr.keras.layers.Bilinear.name_scope",
     "tfr.keras.layers.DocumentInteractionAttention.non_trainable_weights": "tfr.keras.layers.Bilinear.non_trainable_weights",
     "tfr.keras.layers.DocumentInteractionAttention.output": "tfr.keras.layers.Bilinear.output",
+    "tfr.keras.layers.DocumentInteractionAttention.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.layers.DocumentInteractionAttention.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.layers.DocumentInteractionAttention.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.layers.DocumentInteractionAttention.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -371,7 +373,6 @@
     "tfr.keras.layers.FlattenList.__new__": "tfr.keras.layers.Bilinear.__new__",
     "tfr.keras.layers.FlattenList.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.layers.FlattenList.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.layers.FlattenList.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.layers.FlattenList.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.layers.FlattenList.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.layers.FlattenList.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -385,12 +386,14 @@
     "tfr.keras.layers.FlattenList.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.layers.FlattenList.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.layers.FlattenList.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.layers.FlattenList.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.layers.FlattenList.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.layers.FlattenList.metrics": "tfr.keras.layers.Bilinear.metrics",
     "tfr.keras.layers.FlattenList.name": "tfr.keras.layers.Bilinear.name",
     "tfr.keras.layers.FlattenList.name_scope": "tfr.keras.layers.Bilinear.name_scope",
     "tfr.keras.layers.FlattenList.non_trainable_weights": "tfr.keras.layers.Bilinear.non_trainable_weights",
     "tfr.keras.layers.FlattenList.output": "tfr.keras.layers.Bilinear.output",
+    "tfr.keras.layers.FlattenList.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.layers.FlattenList.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.layers.FlattenList.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.layers.FlattenList.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -408,7 +411,6 @@
     "tfr.keras.layers.GAMLayer.__new__": "tfr.keras.layers.Bilinear.__new__",
     "tfr.keras.layers.GAMLayer.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.layers.GAMLayer.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.layers.GAMLayer.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.layers.GAMLayer.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.layers.GAMLayer.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.layers.GAMLayer.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -422,12 +424,14 @@
     "tfr.keras.layers.GAMLayer.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.layers.GAMLayer.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.layers.GAMLayer.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.layers.GAMLayer.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.layers.GAMLayer.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.layers.GAMLayer.metrics": "tfr.keras.layers.Bilinear.metrics",
     "tfr.keras.layers.GAMLayer.name": "tfr.keras.layers.Bilinear.name",
     "tfr.keras.layers.GAMLayer.name_scope": "tfr.keras.layers.Bilinear.name_scope",
     "tfr.keras.layers.GAMLayer.non_trainable_weights": "tfr.keras.layers.Bilinear.non_trainable_weights",
     "tfr.keras.layers.GAMLayer.output": "tfr.keras.layers.Bilinear.output",
+    "tfr.keras.layers.GAMLayer.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.layers.GAMLayer.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.layers.GAMLayer.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.layers.GAMLayer.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -445,7 +449,6 @@
     "tfr.keras.layers.RestoreList.__new__": "tfr.keras.layers.Bilinear.__new__",
     "tfr.keras.layers.RestoreList.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.layers.RestoreList.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.layers.RestoreList.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.layers.RestoreList.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.layers.RestoreList.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.layers.RestoreList.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -459,12 +462,14 @@
     "tfr.keras.layers.RestoreList.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.layers.RestoreList.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.layers.RestoreList.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.layers.RestoreList.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.layers.RestoreList.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.layers.RestoreList.metrics": "tfr.keras.layers.Bilinear.metrics",
     "tfr.keras.layers.RestoreList.name": "tfr.keras.layers.Bilinear.name",
     "tfr.keras.layers.RestoreList.name_scope": "tfr.keras.layers.Bilinear.name_scope",
     "tfr.keras.layers.RestoreList.non_trainable_weights": "tfr.keras.layers.Bilinear.non_trainable_weights",
     "tfr.keras.layers.RestoreList.output": "tfr.keras.layers.Bilinear.output",
+    "tfr.keras.layers.RestoreList.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.layers.RestoreList.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.layers.RestoreList.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.layers.RestoreList.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -482,7 +487,6 @@
     "tfr.keras.layers.SelfAttentionMask.__new__": "tfr.keras.layers.Bilinear.__new__",
     "tfr.keras.layers.SelfAttentionMask.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.layers.SelfAttentionMask.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.layers.SelfAttentionMask.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.layers.SelfAttentionMask.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.layers.SelfAttentionMask.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.layers.SelfAttentionMask.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -496,12 +500,14 @@
     "tfr.keras.layers.SelfAttentionMask.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.layers.SelfAttentionMask.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.layers.SelfAttentionMask.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.layers.SelfAttentionMask.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.layers.SelfAttentionMask.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.layers.SelfAttentionMask.metrics": "tfr.keras.layers.Bilinear.metrics",
     "tfr.keras.layers.SelfAttentionMask.name": "tfr.keras.layers.Bilinear.name",
     "tfr.keras.layers.SelfAttentionMask.name_scope": "tfr.keras.layers.Bilinear.name_scope",
     "tfr.keras.layers.SelfAttentionMask.non_trainable_weights": "tfr.keras.layers.Bilinear.non_trainable_weights",
     "tfr.keras.layers.SelfAttentionMask.output": "tfr.keras.layers.Bilinear.output",
+    "tfr.keras.layers.SelfAttentionMask.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.layers.SelfAttentionMask.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.layers.SelfAttentionMask.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.layers.SelfAttentionMask.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -697,6 +703,21 @@
     "tfr.keras.losses.UniqueSoftmaxLoss.__ne__": "tfr.keras.layers.Bilinear.__ne__",
     "tfr.keras.losses.UniqueSoftmaxLoss.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
     "tfr.keras.losses.UniqueSoftmaxLoss.get_config": "tfr.keras.losses.ApproxMRRLoss.get_config",
+    "tfr.keras.losses.YetiDCGLambdaWeight.__eq__": "tfr.keras.layers.Bilinear.__eq__",
+    "tfr.keras.losses.YetiDCGLambdaWeight.__ge__": "tfr.keras.layers.Bilinear.__ge__",
+    "tfr.keras.losses.YetiDCGLambdaWeight.__gt__": "tfr.keras.layers.Bilinear.__gt__",
+    "tfr.keras.losses.YetiDCGLambdaWeight.__le__": "tfr.keras.layers.Bilinear.__le__",
+    "tfr.keras.losses.YetiDCGLambdaWeight.__lt__": "tfr.keras.layers.Bilinear.__lt__",
+    "tfr.keras.losses.YetiDCGLambdaWeight.__ne__": "tfr.keras.layers.Bilinear.__ne__",
+    "tfr.keras.losses.YetiDCGLambdaWeight.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
+    "tfr.keras.losses.YetiDCGLambdaWeight.individual_weights": "tfr.keras.losses.DCGLambdaWeight.individual_weights",
+    "tfr.keras.losses.YetiLogisticLoss.__eq__": "tfr.keras.layers.Bilinear.__eq__",
+    "tfr.keras.losses.YetiLogisticLoss.__ge__": "tfr.keras.layers.Bilinear.__ge__",
+    "tfr.keras.losses.YetiLogisticLoss.__gt__": "tfr.keras.layers.Bilinear.__gt__",
+    "tfr.keras.losses.YetiLogisticLoss.__le__": "tfr.keras.layers.Bilinear.__le__",
+    "tfr.keras.losses.YetiLogisticLoss.__lt__": "tfr.keras.layers.Bilinear.__lt__",
+    "tfr.keras.losses.YetiLogisticLoss.__ne__": "tfr.keras.layers.Bilinear.__ne__",
+    "tfr.keras.losses.YetiLogisticLoss.__new__": "tfr.keras.losses.ApproxMRRLoss.__new__",
     "tfr.keras.metrics.ARPMetric.__eq__": "tfr.keras.layers.Bilinear.__eq__",
     "tfr.keras.metrics.ARPMetric.__ge__": "tfr.keras.layers.Bilinear.__ge__",
     "tfr.keras.metrics.ARPMetric.__gt__": "tfr.keras.layers.Bilinear.__gt__",
@@ -705,7 +726,6 @@
     "tfr.keras.metrics.ARPMetric.__ne__": "tfr.keras.layers.Bilinear.__ne__",
     "tfr.keras.metrics.ARPMetric.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.metrics.ARPMetric.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.metrics.ARPMetric.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.metrics.ARPMetric.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.metrics.ARPMetric.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.metrics.ARPMetric.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -718,11 +738,13 @@
     "tfr.keras.metrics.ARPMetric.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.metrics.ARPMetric.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.metrics.ARPMetric.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.metrics.ARPMetric.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.metrics.ARPMetric.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.metrics.ARPMetric.metrics": "tfr.keras.layers.Bilinear.metrics",
     "tfr.keras.metrics.ARPMetric.name": "tfr.keras.layers.Bilinear.name",
     "tfr.keras.metrics.ARPMetric.name_scope": "tfr.keras.layers.Bilinear.name_scope",
     "tfr.keras.metrics.ARPMetric.output": "tfr.keras.layers.Bilinear.output",
+    "tfr.keras.metrics.ARPMetric.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.metrics.ARPMetric.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.metrics.ARPMetric.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.metrics.ARPMetric.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -739,7 +761,6 @@
     "tfr.keras.metrics.AlphaDCGMetric.__new__": "tfr.keras.metrics.ARPMetric.__new__",
     "tfr.keras.metrics.AlphaDCGMetric.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.metrics.AlphaDCGMetric.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.metrics.AlphaDCGMetric.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.metrics.AlphaDCGMetric.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.metrics.AlphaDCGMetric.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.metrics.AlphaDCGMetric.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -753,6 +774,7 @@
     "tfr.keras.metrics.AlphaDCGMetric.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.metrics.AlphaDCGMetric.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.metrics.AlphaDCGMetric.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.metrics.AlphaDCGMetric.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.metrics.AlphaDCGMetric.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.metrics.AlphaDCGMetric.merge_state": "tfr.keras.metrics.ARPMetric.merge_state",
     "tfr.keras.metrics.AlphaDCGMetric.metrics": "tfr.keras.layers.Bilinear.metrics",
@@ -762,6 +784,7 @@
     "tfr.keras.metrics.AlphaDCGMetric.output": "tfr.keras.layers.Bilinear.output",
     "tfr.keras.metrics.AlphaDCGMetric.reset_state": "tfr.keras.metrics.ARPMetric.reset_state",
     "tfr.keras.metrics.AlphaDCGMetric.result": "tfr.keras.metrics.ARPMetric.result",
+    "tfr.keras.metrics.AlphaDCGMetric.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.metrics.AlphaDCGMetric.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.metrics.AlphaDCGMetric.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.metrics.AlphaDCGMetric.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -780,7 +803,6 @@
     "tfr.keras.metrics.DCGMetric.__new__": "tfr.keras.metrics.ARPMetric.__new__",
     "tfr.keras.metrics.DCGMetric.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.metrics.DCGMetric.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.metrics.DCGMetric.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.metrics.DCGMetric.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.metrics.DCGMetric.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.metrics.DCGMetric.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -794,6 +816,7 @@
     "tfr.keras.metrics.DCGMetric.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.metrics.DCGMetric.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.metrics.DCGMetric.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.metrics.DCGMetric.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.metrics.DCGMetric.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.metrics.DCGMetric.merge_state": "tfr.keras.metrics.ARPMetric.merge_state",
     "tfr.keras.metrics.DCGMetric.metrics": "tfr.keras.layers.Bilinear.metrics",
@@ -803,6 +826,7 @@
     "tfr.keras.metrics.DCGMetric.output": "tfr.keras.layers.Bilinear.output",
     "tfr.keras.metrics.DCGMetric.reset_state": "tfr.keras.metrics.ARPMetric.reset_state",
     "tfr.keras.metrics.DCGMetric.result": "tfr.keras.metrics.ARPMetric.result",
+    "tfr.keras.metrics.DCGMetric.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.metrics.DCGMetric.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.metrics.DCGMetric.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.metrics.DCGMetric.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -821,7 +845,6 @@
     "tfr.keras.metrics.HitsMetric.__new__": "tfr.keras.metrics.ARPMetric.__new__",
     "tfr.keras.metrics.HitsMetric.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.metrics.HitsMetric.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.metrics.HitsMetric.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.metrics.HitsMetric.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.metrics.HitsMetric.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.metrics.HitsMetric.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -835,6 +858,7 @@
     "tfr.keras.metrics.HitsMetric.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.metrics.HitsMetric.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.metrics.HitsMetric.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.metrics.HitsMetric.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.metrics.HitsMetric.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.metrics.HitsMetric.merge_state": "tfr.keras.metrics.ARPMetric.merge_state",
     "tfr.keras.metrics.HitsMetric.metrics": "tfr.keras.layers.Bilinear.metrics",
@@ -844,6 +868,7 @@
     "tfr.keras.metrics.HitsMetric.output": "tfr.keras.layers.Bilinear.output",
     "tfr.keras.metrics.HitsMetric.reset_state": "tfr.keras.metrics.ARPMetric.reset_state",
     "tfr.keras.metrics.HitsMetric.result": "tfr.keras.metrics.ARPMetric.result",
+    "tfr.keras.metrics.HitsMetric.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.metrics.HitsMetric.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.metrics.HitsMetric.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.metrics.HitsMetric.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -862,7 +887,6 @@
     "tfr.keras.metrics.MRRMetric.__new__": "tfr.keras.metrics.ARPMetric.__new__",
     "tfr.keras.metrics.MRRMetric.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.metrics.MRRMetric.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.metrics.MRRMetric.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.metrics.MRRMetric.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.metrics.MRRMetric.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.metrics.MRRMetric.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -876,6 +900,7 @@
     "tfr.keras.metrics.MRRMetric.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.metrics.MRRMetric.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.metrics.MRRMetric.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.metrics.MRRMetric.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.metrics.MRRMetric.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.metrics.MRRMetric.merge_state": "tfr.keras.metrics.ARPMetric.merge_state",
     "tfr.keras.metrics.MRRMetric.metrics": "tfr.keras.layers.Bilinear.metrics",
@@ -885,6 +910,7 @@
     "tfr.keras.metrics.MRRMetric.output": "tfr.keras.layers.Bilinear.output",
     "tfr.keras.metrics.MRRMetric.reset_state": "tfr.keras.metrics.ARPMetric.reset_state",
     "tfr.keras.metrics.MRRMetric.result": "tfr.keras.metrics.ARPMetric.result",
+    "tfr.keras.metrics.MRRMetric.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.metrics.MRRMetric.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.metrics.MRRMetric.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.metrics.MRRMetric.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -903,7 +929,6 @@
     "tfr.keras.metrics.MeanAveragePrecisionMetric.__new__": "tfr.keras.metrics.ARPMetric.__new__",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.metrics.MeanAveragePrecisionMetric.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -917,6 +942,7 @@
     "tfr.keras.metrics.MeanAveragePrecisionMetric.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.metrics.MeanAveragePrecisionMetric.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.merge_state": "tfr.keras.metrics.ARPMetric.merge_state",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.metrics": "tfr.keras.layers.Bilinear.metrics",
@@ -926,6 +952,7 @@
     "tfr.keras.metrics.MeanAveragePrecisionMetric.output": "tfr.keras.layers.Bilinear.output",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.reset_state": "tfr.keras.metrics.ARPMetric.reset_state",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.result": "tfr.keras.metrics.ARPMetric.result",
+    "tfr.keras.metrics.MeanAveragePrecisionMetric.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.metrics.MeanAveragePrecisionMetric.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -944,7 +971,6 @@
     "tfr.keras.metrics.NDCGMetric.__new__": "tfr.keras.metrics.ARPMetric.__new__",
     "tfr.keras.metrics.NDCGMetric.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.metrics.NDCGMetric.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.metrics.NDCGMetric.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.metrics.NDCGMetric.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.metrics.NDCGMetric.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.metrics.NDCGMetric.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -958,6 +984,7 @@
     "tfr.keras.metrics.NDCGMetric.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.metrics.NDCGMetric.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.metrics.NDCGMetric.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.metrics.NDCGMetric.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.metrics.NDCGMetric.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.metrics.NDCGMetric.merge_state": "tfr.keras.metrics.ARPMetric.merge_state",
     "tfr.keras.metrics.NDCGMetric.metrics": "tfr.keras.layers.Bilinear.metrics",
@@ -967,6 +994,7 @@
     "tfr.keras.metrics.NDCGMetric.output": "tfr.keras.layers.Bilinear.output",
     "tfr.keras.metrics.NDCGMetric.reset_state": "tfr.keras.metrics.ARPMetric.reset_state",
     "tfr.keras.metrics.NDCGMetric.result": "tfr.keras.metrics.ARPMetric.result",
+    "tfr.keras.metrics.NDCGMetric.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.metrics.NDCGMetric.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.metrics.NDCGMetric.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.metrics.NDCGMetric.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -985,7 +1013,6 @@
     "tfr.keras.metrics.OPAMetric.__new__": "tfr.keras.metrics.ARPMetric.__new__",
     "tfr.keras.metrics.OPAMetric.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.metrics.OPAMetric.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.metrics.OPAMetric.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.metrics.OPAMetric.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.metrics.OPAMetric.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.metrics.OPAMetric.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -1000,6 +1027,7 @@
     "tfr.keras.metrics.OPAMetric.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.metrics.OPAMetric.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.metrics.OPAMetric.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.metrics.OPAMetric.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.metrics.OPAMetric.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.metrics.OPAMetric.merge_state": "tfr.keras.metrics.ARPMetric.merge_state",
     "tfr.keras.metrics.OPAMetric.metrics": "tfr.keras.layers.Bilinear.metrics",
@@ -1009,6 +1037,7 @@
     "tfr.keras.metrics.OPAMetric.output": "tfr.keras.layers.Bilinear.output",
     "tfr.keras.metrics.OPAMetric.reset_state": "tfr.keras.metrics.ARPMetric.reset_state",
     "tfr.keras.metrics.OPAMetric.result": "tfr.keras.metrics.ARPMetric.result",
+    "tfr.keras.metrics.OPAMetric.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.metrics.OPAMetric.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.metrics.OPAMetric.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.metrics.OPAMetric.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -1027,7 +1056,6 @@
     "tfr.keras.metrics.PrecisionIAMetric.__new__": "tfr.keras.metrics.ARPMetric.__new__",
     "tfr.keras.metrics.PrecisionIAMetric.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.metrics.PrecisionIAMetric.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.metrics.PrecisionIAMetric.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.metrics.PrecisionIAMetric.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.metrics.PrecisionIAMetric.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.metrics.PrecisionIAMetric.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -1041,6 +1069,7 @@
     "tfr.keras.metrics.PrecisionIAMetric.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.metrics.PrecisionIAMetric.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.metrics.PrecisionIAMetric.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.metrics.PrecisionIAMetric.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.metrics.PrecisionIAMetric.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.metrics.PrecisionIAMetric.merge_state": "tfr.keras.metrics.ARPMetric.merge_state",
     "tfr.keras.metrics.PrecisionIAMetric.metrics": "tfr.keras.layers.Bilinear.metrics",
@@ -1050,6 +1079,7 @@
     "tfr.keras.metrics.PrecisionIAMetric.output": "tfr.keras.layers.Bilinear.output",
     "tfr.keras.metrics.PrecisionIAMetric.reset_state": "tfr.keras.metrics.ARPMetric.reset_state",
     "tfr.keras.metrics.PrecisionIAMetric.result": "tfr.keras.metrics.ARPMetric.result",
+    "tfr.keras.metrics.PrecisionIAMetric.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.metrics.PrecisionIAMetric.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.metrics.PrecisionIAMetric.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.metrics.PrecisionIAMetric.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -1068,7 +1098,6 @@
     "tfr.keras.metrics.PrecisionMetric.__new__": "tfr.keras.metrics.ARPMetric.__new__",
     "tfr.keras.metrics.PrecisionMetric.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.metrics.PrecisionMetric.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.metrics.PrecisionMetric.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.metrics.PrecisionMetric.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.metrics.PrecisionMetric.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.metrics.PrecisionMetric.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -1082,6 +1111,7 @@
     "tfr.keras.metrics.PrecisionMetric.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.metrics.PrecisionMetric.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.metrics.PrecisionMetric.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.metrics.PrecisionMetric.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.metrics.PrecisionMetric.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.metrics.PrecisionMetric.merge_state": "tfr.keras.metrics.ARPMetric.merge_state",
     "tfr.keras.metrics.PrecisionMetric.metrics": "tfr.keras.layers.Bilinear.metrics",
@@ -1091,6 +1121,7 @@
     "tfr.keras.metrics.PrecisionMetric.output": "tfr.keras.layers.Bilinear.output",
     "tfr.keras.metrics.PrecisionMetric.reset_state": "tfr.keras.metrics.ARPMetric.reset_state",
     "tfr.keras.metrics.PrecisionMetric.result": "tfr.keras.metrics.ARPMetric.result",
+    "tfr.keras.metrics.PrecisionMetric.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.metrics.PrecisionMetric.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.metrics.PrecisionMetric.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.metrics.PrecisionMetric.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -1117,7 +1148,6 @@
     "tfr.keras.metrics.RecallMetric.__new__": "tfr.keras.metrics.ARPMetric.__new__",
     "tfr.keras.metrics.RecallMetric.activity_regularizer": "tfr.keras.layers.Bilinear.activity_regularizer",
     "tfr.keras.metrics.RecallMetric.add_loss": "tfr.keras.layers.Bilinear.add_loss",
-    "tfr.keras.metrics.RecallMetric.add_metric": "tfr.keras.layers.Bilinear.add_metric",
     "tfr.keras.metrics.RecallMetric.build": "tfr.keras.layers.ConcatFeatures.build",
     "tfr.keras.metrics.RecallMetric.build_from_config": "tfr.keras.layers.Bilinear.build_from_config",
     "tfr.keras.metrics.RecallMetric.compute_dtype": "tfr.keras.layers.Bilinear.compute_dtype",
@@ -1131,6 +1161,7 @@
     "tfr.keras.metrics.RecallMetric.get_weights": "tfr.keras.layers.Bilinear.get_weights",
     "tfr.keras.metrics.RecallMetric.input": "tfr.keras.layers.Bilinear.input",
     "tfr.keras.metrics.RecallMetric.input_spec": "tfr.keras.layers.Bilinear.input_spec",
+    "tfr.keras.metrics.RecallMetric.load_own_variables": "tfr.keras.layers.Bilinear.load_own_variables",
     "tfr.keras.metrics.RecallMetric.losses": "tfr.keras.layers.Bilinear.losses",
     "tfr.keras.metrics.RecallMetric.merge_state": "tfr.keras.metrics.ARPMetric.merge_state",
     "tfr.keras.metrics.RecallMetric.metrics": "tfr.keras.layers.Bilinear.metrics",
@@ -1140,6 +1171,7 @@
     "tfr.keras.metrics.RecallMetric.output": "tfr.keras.layers.Bilinear.output",
     "tfr.keras.metrics.RecallMetric.reset_state": "tfr.keras.metrics.ARPMetric.reset_state",
     "tfr.keras.metrics.RecallMetric.result": "tfr.keras.metrics.ARPMetric.result",
+    "tfr.keras.metrics.RecallMetric.save_own_variables": "tfr.keras.layers.Bilinear.save_own_variables",
     "tfr.keras.metrics.RecallMetric.set_weights": "tfr.keras.layers.Bilinear.set_weights",
     "tfr.keras.metrics.RecallMetric.submodules": "tfr.keras.layers.Bilinear.submodules",
     "tfr.keras.metrics.RecallMetric.supports_masking": "tfr.keras.layers.Bilinear.supports_masking",
@@ -1418,6 +1450,7 @@
     "tfr.extension.premade.TFRBertConfig.__ne__": true,
     "tfr.extension.premade.TFRBertConfig.__new__": true,
     "tfr.extension.premade.TFRBertConfig.aggregated_metrics": true,
+    "tfr.extension.premade.TFRBertConfig.allow_image_summary": true,
     "tfr.extension.premade.TFRBertConfig.as_dict": true,
     "tfr.extension.premade.TFRBertConfig.default_params": true,
     "tfr.extension.premade.TFRBertConfig.differential_privacy_config": true,
@@ -1454,6 +1487,7 @@
     "tfr.extension.premade.TFRBertDataConfig.__new__": true,
     "tfr.extension.premade.TFRBertDataConfig.apply_tf_data_service_before_batching": true,
     "tfr.extension.premade.TFRBertDataConfig.as_dict": true,
+    "tfr.extension.premade.TFRBertDataConfig.autotune_algorithm": true,
     "tfr.extension.premade.TFRBertDataConfig.block_length": true,
     "tfr.extension.premade.TFRBertDataConfig.cache": true,
     "tfr.extension.premade.TFRBertDataConfig.convert_labels_to_binary": true,
@@ -1607,6 +1641,7 @@
     "tfr.extension.premade.tfrbert_task.TFRBertConfig.__ne__": true,
     "tfr.extension.premade.tfrbert_task.TFRBertConfig.__new__": true,
     "tfr.extension.premade.tfrbert_task.TFRBertConfig.aggregated_metrics": true,
+    "tfr.extension.premade.tfrbert_task.TFRBertConfig.allow_image_summary": true,
     "tfr.extension.premade.tfrbert_task.TFRBertConfig.as_dict": true,
     "tfr.extension.premade.tfrbert_task.TFRBertConfig.default_params": true,
     "tfr.extension.premade.tfrbert_task.TFRBertConfig.differential_privacy_config": true,
@@ -1643,6 +1678,7 @@
     "tfr.extension.premade.tfrbert_task.TFRBertDataConfig.__new__": true,
     "tfr.extension.premade.tfrbert_task.TFRBertDataConfig.apply_tf_data_service_before_batching": true,
     "tfr.extension.premade.tfrbert_task.TFRBertDataConfig.as_dict": true,
+    "tfr.extension.premade.tfrbert_task.TFRBertDataConfig.autotune_algorithm": true,
     "tfr.extension.premade.tfrbert_task.TFRBertDataConfig.block_length": true,
     "tfr.extension.premade.tfrbert_task.TFRBertDataConfig.cache": true,
     "tfr.extension.premade.tfrbert_task.TFRBertDataConfig.convert_labels_to_binary": true,
@@ -1798,6 +1834,7 @@
     "tfr.extension.task.RankingDataConfig.__new__": true,
     "tfr.extension.task.RankingDataConfig.apply_tf_data_service_before_batching": true,
     "tfr.extension.task.RankingDataConfig.as_dict": true,
+    "tfr.extension.task.RankingDataConfig.autotune_algorithm": true,
     "tfr.extension.task.RankingDataConfig.block_length": true,
     "tfr.extension.task.RankingDataConfig.cache": true,
     "tfr.extension.task.RankingDataConfig.convert_labels_to_binary": true,
@@ -1893,6 +1930,7 @@
     "tfr.extension.task.RankingTaskConfig.__ne__": true,
     "tfr.extension.task.RankingTaskConfig.__new__": true,
     "tfr.extension.task.RankingTaskConfig.aggregated_metrics": true,
+    "tfr.extension.task.RankingTaskConfig.allow_image_summary": true,
     "tfr.extension.task.RankingTaskConfig.as_dict": true,
     "tfr.extension.task.RankingTaskConfig.default_params": true,
     "tfr.extension.task.RankingTaskConfig.differential_privacy_config": true,
@@ -1927,7 +1965,6 @@
     "tfr.keras.layers.Bilinear.__new__": true,
     "tfr.keras.layers.Bilinear.activity_regularizer": true,
     "tfr.keras.layers.Bilinear.add_loss": true,
-    "tfr.keras.layers.Bilinear.add_metric": true,
     "tfr.keras.layers.Bilinear.build": true,
     "tfr.keras.layers.Bilinear.build_from_config": true,
     "tfr.keras.layers.Bilinear.compute_dtype": true,
@@ -1943,12 +1980,14 @@
     "tfr.keras.layers.Bilinear.get_weights": true,
     "tfr.keras.layers.Bilinear.input": true,
     "tfr.keras.layers.Bilinear.input_spec": true,
+    "tfr.keras.layers.Bilinear.load_own_variables": true,
     "tfr.keras.layers.Bilinear.losses": true,
     "tfr.keras.layers.Bilinear.metrics": true,
     "tfr.keras.layers.Bilinear.name": true,
     "tfr.keras.layers.Bilinear.name_scope": true,
     "tfr.keras.layers.Bilinear.non_trainable_weights": true,
     "tfr.keras.layers.Bilinear.output": true,
+    "tfr.keras.layers.Bilinear.save_own_variables": true,
     "tfr.keras.layers.Bilinear.set_weights": true,
     "tfr.keras.layers.Bilinear.submodules": true,
     "tfr.keras.layers.Bilinear.supports_masking": true,
@@ -1969,7 +2008,6 @@
     "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.build": true,
     "tfr.keras.layers.ConcatFeatures.build_from_config": true,
     "tfr.keras.layers.ConcatFeatures.compute_dtype": true,
@@ -1985,12 +2023,14 @@
     "tfr.keras.layers.ConcatFeatures.get_weights": true,
     "tfr.keras.layers.ConcatFeatures.input": true,
     "tfr.keras.layers.ConcatFeatures.input_spec": true,
+    "tfr.keras.layers.ConcatFeatures.load_own_variables": 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.save_own_variables": true,
     "tfr.keras.layers.ConcatFeatures.set_weights": true,
     "tfr.keras.layers.ConcatFeatures.submodules": true,
     "tfr.keras.layers.ConcatFeatures.supports_masking": true,
@@ -2011,7 +2051,6 @@
     "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.build": true,
     "tfr.keras.layers.DocumentInteractionAttention.build_from_config": true,
     "tfr.keras.layers.DocumentInteractionAttention.compute_dtype": true,
@@ -2027,12 +2066,14 @@
     "tfr.keras.layers.DocumentInteractionAttention.get_weights": true,
     "tfr.keras.layers.DocumentInteractionAttention.input": true,
     "tfr.keras.layers.DocumentInteractionAttention.input_spec": true,
+    "tfr.keras.layers.DocumentInteractionAttention.load_own_variables": 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.save_own_variables": true,
     "tfr.keras.layers.DocumentInteractionAttention.set_weights": true,
     "tfr.keras.layers.DocumentInteractionAttention.submodules": true,
     "tfr.keras.layers.DocumentInteractionAttention.supports_masking": true,
@@ -2053,7 +2094,6 @@
     "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.build": true,
     "tfr.keras.layers.FlattenList.build_from_config": true,
     "tfr.keras.layers.FlattenList.compute_dtype": true,
@@ -2069,12 +2109,14 @@
     "tfr.keras.layers.FlattenList.get_weights": true,
     "tfr.keras.layers.FlattenList.input": true,
     "tfr.keras.layers.FlattenList.input_spec": true,
+    "tfr.keras.layers.FlattenList.load_own_variables": 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.save_own_variables": true,
     "tfr.keras.layers.FlattenList.set_weights": true,
     "tfr.keras.layers.FlattenList.submodules": true,
     "tfr.keras.layers.FlattenList.supports_masking": true,
@@ -2095,7 +2137,6 @@
     "tfr.keras.layers.GAMLayer.__new__": true,
     "tfr.keras.layers.GAMLayer.activity_regularizer": true,
     "tfr.keras.layers.GAMLayer.add_loss": true,
-    "tfr.keras.layers.GAMLayer.add_metric": true,
     "tfr.keras.layers.GAMLayer.build": true,
     "tfr.keras.layers.GAMLayer.build_from_config": true,
     "tfr.keras.layers.GAMLayer.compute_dtype": true,
@@ -2111,12 +2152,14 @@
     "tfr.keras.layers.GAMLayer.get_weights": true,
     "tfr.keras.layers.GAMLayer.input": true,
     "tfr.keras.layers.GAMLayer.input_spec": true,
+    "tfr.keras.layers.GAMLayer.load_own_variables": true,
     "tfr.keras.layers.GAMLayer.losses": true,
     "tfr.keras.layers.GAMLayer.metrics": true,
     "tfr.keras.layers.GAMLayer.name": true,
     "tfr.keras.layers.GAMLayer.name_scope": true,
     "tfr.keras.layers.GAMLayer.non_trainable_weights": true,
     "tfr.keras.layers.GAMLayer.output": true,
+    "tfr.keras.layers.GAMLayer.save_own_variables": true,
     "tfr.keras.layers.GAMLayer.set_weights": true,
     "tfr.keras.layers.GAMLayer.submodules": true,
     "tfr.keras.layers.GAMLayer.supports_masking": true,
@@ -2137,7 +2180,6 @@
     "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.build": true,
     "tfr.keras.layers.RestoreList.build_from_config": true,
     "tfr.keras.layers.RestoreList.compute_dtype": true,
@@ -2153,12 +2195,14 @@
     "tfr.keras.layers.RestoreList.get_weights": true,
     "tfr.keras.layers.RestoreList.input": true,
     "tfr.keras.layers.RestoreList.input_spec": true,
+    "tfr.keras.layers.RestoreList.load_own_variables": 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.save_own_variables": true,
     "tfr.keras.layers.RestoreList.set_weights": true,
     "tfr.keras.layers.RestoreList.submodules": true,
     "tfr.keras.layers.RestoreList.supports_masking": true,
@@ -2179,7 +2223,6 @@
     "tfr.keras.layers.SelfAttentionMask.__new__": true,
     "tfr.keras.layers.SelfAttentionMask.activity_regularizer": true,
     "tfr.keras.layers.SelfAttentionMask.add_loss": true,
-    "tfr.keras.layers.SelfAttentionMask.add_metric": true,
     "tfr.keras.layers.SelfAttentionMask.build": true,
     "tfr.keras.layers.SelfAttentionMask.build_from_config": true,
     "tfr.keras.layers.SelfAttentionMask.compute_dtype": true,
@@ -2195,12 +2238,14 @@
     "tfr.keras.layers.SelfAttentionMask.get_weights": true,
     "tfr.keras.layers.SelfAttentionMask.input": true,
     "tfr.keras.layers.SelfAttentionMask.input_spec": true,
+    "tfr.keras.layers.SelfAttentionMask.load_own_variables": true,
     "tfr.keras.layers.SelfAttentionMask.losses": true,
     "tfr.keras.layers.SelfAttentionMask.metrics": true,
     "tfr.keras.layers.SelfAttentionMask.name": true,
     "tfr.keras.layers.SelfAttentionMask.name_scope": true,
     "tfr.keras.layers.SelfAttentionMask.non_trainable_weights": true,
     "tfr.keras.layers.SelfAttentionMask.output": true,
+    "tfr.keras.layers.SelfAttentionMask.save_own_variables": true,
     "tfr.keras.layers.SelfAttentionMask.set_weights": true,
     "tfr.keras.layers.SelfAttentionMask.submodules": true,
     "tfr.keras.layers.SelfAttentionMask.supports_masking": true,
@@ -2454,6 +2499,7 @@
     "tfr.keras.losses.RankingLossKey.SIGMOID_CROSS_ENTROPY_LOSS": true,
     "tfr.keras.losses.RankingLossKey.SOFTMAX_LOSS": true,
     "tfr.keras.losses.RankingLossKey.UNIQUE_SOFTMAX_LOSS": true,
+    "tfr.keras.losses.RankingLossKey.YETI_LOGISTIC_LOSS": true,
     "tfr.keras.losses.RankingLossKey.__eq__": true,
     "tfr.keras.losses.RankingLossKey.__ge__": true,
     "tfr.keras.losses.RankingLossKey.__gt__": true,
@@ -2499,6 +2545,30 @@
     "tfr.keras.losses.UniqueSoftmaxLoss.__new__": true,
     "tfr.keras.losses.UniqueSoftmaxLoss.from_config": true,
     "tfr.keras.losses.UniqueSoftmaxLoss.get_config": true,
+    "tfr.keras.losses.YetiDCGLambdaWeight": false,
+    "tfr.keras.losses.YetiDCGLambdaWeight.__eq__": true,
+    "tfr.keras.losses.YetiDCGLambdaWeight.__ge__": true,
+    "tfr.keras.losses.YetiDCGLambdaWeight.__gt__": true,
+    "tfr.keras.losses.YetiDCGLambdaWeight.__init__": true,
+    "tfr.keras.losses.YetiDCGLambdaWeight.__le__": true,
+    "tfr.keras.losses.YetiDCGLambdaWeight.__lt__": true,
+    "tfr.keras.losses.YetiDCGLambdaWeight.__ne__": true,
+    "tfr.keras.losses.YetiDCGLambdaWeight.__new__": true,
+    "tfr.keras.losses.YetiDCGLambdaWeight.get_config": true,
+    "tfr.keras.losses.YetiDCGLambdaWeight.individual_weights": true,
+    "tfr.keras.losses.YetiDCGLambdaWeight.pair_weights": true,
+    "tfr.keras.losses.YetiLogisticLoss": false,
+    "tfr.keras.losses.YetiLogisticLoss.__call__": true,
+    "tfr.keras.losses.YetiLogisticLoss.__eq__": true,
+    "tfr.keras.losses.YetiLogisticLoss.__ge__": true,
+    "tfr.keras.losses.YetiLogisticLoss.__gt__": true,
+    "tfr.keras.losses.YetiLogisticLoss.__init__": true,
+    "tfr.keras.losses.YetiLogisticLoss.__le__": true,
+    "tfr.keras.losses.YetiLogisticLoss.__lt__": true,
+    "tfr.keras.losses.YetiLogisticLoss.__ne__": true,
+    "tfr.keras.losses.YetiLogisticLoss.__new__": true,
+    "tfr.keras.losses.YetiLogisticLoss.from_config": true,
+    "tfr.keras.losses.YetiLogisticLoss.get_config": true,
     "tfr.keras.losses.get": false,
     "tfr.keras.metrics": false,
     "tfr.keras.metrics.ARPMetric": false,
@@ -2513,7 +2583,6 @@
     "tfr.keras.metrics.ARPMetric.__new__": true,
     "tfr.keras.metrics.ARPMetric.activity_regularizer": true,
     "tfr.keras.metrics.ARPMetric.add_loss": true,
-    "tfr.keras.metrics.ARPMetric.add_metric": true,
     "tfr.keras.metrics.ARPMetric.build": true,
     "tfr.keras.metrics.ARPMetric.build_from_config": true,
     "tfr.keras.metrics.ARPMetric.compute_dtype": true,
@@ -2529,6 +2598,7 @@
     "tfr.keras.metrics.ARPMetric.get_weights": true,
     "tfr.keras.metrics.ARPMetric.input": true,
     "tfr.keras.metrics.ARPMetric.input_spec": true,
+    "tfr.keras.metrics.ARPMetric.load_own_variables": true,
     "tfr.keras.metrics.ARPMetric.losses": true,
     "tfr.keras.metrics.ARPMetric.merge_state": true,
     "tfr.keras.metrics.ARPMetric.metrics": true,
@@ -2538,6 +2608,7 @@
     "tfr.keras.metrics.ARPMetric.output": true,
     "tfr.keras.metrics.ARPMetric.reset_state": true,
     "tfr.keras.metrics.ARPMetric.result": true,
+    "tfr.keras.metrics.ARPMetric.save_own_variables": true,
     "tfr.keras.metrics.ARPMetric.set_weights": true,
     "tfr.keras.metrics.ARPMetric.submodules": true,
     "tfr.keras.metrics.ARPMetric.supports_masking": true,
@@ -2559,7 +2630,6 @@
     "tfr.keras.metrics.AlphaDCGMetric.__new__": true,
     "tfr.keras.metrics.AlphaDCGMetric.activity_regularizer": true,
     "tfr.keras.metrics.AlphaDCGMetric.add_loss": true,
-    "tfr.keras.metrics.AlphaDCGMetric.add_metric": true,
     "tfr.keras.metrics.AlphaDCGMetric.build": true,
     "tfr.keras.metrics.AlphaDCGMetric.build_from_config": true,
     "tfr.keras.metrics.AlphaDCGMetric.compute_dtype": true,
@@ -2575,6 +2645,7 @@
     "tfr.keras.metrics.AlphaDCGMetric.get_weights": true,
     "tfr.keras.metrics.AlphaDCGMetric.input": true,
     "tfr.keras.metrics.AlphaDCGMetric.input_spec": true,
+    "tfr.keras.metrics.AlphaDCGMetric.load_own_variables": true,
     "tfr.keras.metrics.AlphaDCGMetric.losses": true,
     "tfr.keras.metrics.AlphaDCGMetric.merge_state": true,
     "tfr.keras.metrics.AlphaDCGMetric.metrics": true,
@@ -2584,6 +2655,7 @@
     "tfr.keras.metrics.AlphaDCGMetric.output": true,
     "tfr.keras.metrics.AlphaDCGMetric.reset_state": true,
     "tfr.keras.metrics.AlphaDCGMetric.result": true,
+    "tfr.keras.metrics.AlphaDCGMetric.save_own_variables": true,
     "tfr.keras.metrics.AlphaDCGMetric.set_weights": true,
     "tfr.keras.metrics.AlphaDCGMetric.submodules": true,
     "tfr.keras.metrics.AlphaDCGMetric.supports_masking": true,
@@ -2605,7 +2677,6 @@
     "tfr.keras.metrics.DCGMetric.__new__": true,
     "tfr.keras.metrics.DCGMetric.activity_regularizer": true,
     "tfr.keras.metrics.DCGMetric.add_loss": true,
-    "tfr.keras.metrics.DCGMetric.add_metric": true,
     "tfr.keras.metrics.DCGMetric.build": true,
     "tfr.keras.metrics.DCGMetric.build_from_config": true,
     "tfr.keras.metrics.DCGMetric.compute_dtype": true,
@@ -2621,6 +2692,7 @@
     "tfr.keras.metrics.DCGMetric.get_weights": true,
     "tfr.keras.metrics.DCGMetric.input": true,
     "tfr.keras.metrics.DCGMetric.input_spec": true,
+    "tfr.keras.metrics.DCGMetric.load_own_variables": true,
     "tfr.keras.metrics.DCGMetric.losses": true,
     "tfr.keras.metrics.DCGMetric.merge_state": true,
     "tfr.keras.metrics.DCGMetric.metrics": true,
@@ -2630,6 +2702,7 @@
     "tfr.keras.metrics.DCGMetric.output": true,
     "tfr.keras.metrics.DCGMetric.reset_state": true,
     "tfr.keras.metrics.DCGMetric.result": true,
+    "tfr.keras.metrics.DCGMetric.save_own_variables": true,
     "tfr.keras.metrics.DCGMetric.set_weights": true,
     "tfr.keras.metrics.DCGMetric.submodules": true,
     "tfr.keras.metrics.DCGMetric.supports_masking": true,
@@ -2651,7 +2724,6 @@
     "tfr.keras.metrics.HitsMetric.__new__": true,
     "tfr.keras.metrics.HitsMetric.activity_regularizer": true,
     "tfr.keras.metrics.HitsMetric.add_loss": true,
-    "tfr.keras.metrics.HitsMetric.add_metric": true,
     "tfr.keras.metrics.HitsMetric.build": true,
     "tfr.keras.metrics.HitsMetric.build_from_config": true,
     "tfr.keras.metrics.HitsMetric.compute_dtype": true,
@@ -2667,6 +2739,7 @@
     "tfr.keras.metrics.HitsMetric.get_weights": true,
     "tfr.keras.metrics.HitsMetric.input": true,
     "tfr.keras.metrics.HitsMetric.input_spec": true,
+    "tfr.keras.metrics.HitsMetric.load_own_variables": true,
     "tfr.keras.metrics.HitsMetric.losses": true,
     "tfr.keras.metrics.HitsMetric.merge_state": true,
     "tfr.keras.metrics.HitsMetric.metrics": true,
@@ -2676,6 +2749,7 @@
     "tfr.keras.metrics.HitsMetric.output": true,
     "tfr.keras.metrics.HitsMetric.reset_state": true,
     "tfr.keras.metrics.HitsMetric.result": true,
+    "tfr.keras.metrics.HitsMetric.save_own_variables": true,
     "tfr.keras.metrics.HitsMetric.set_weights": true,
     "tfr.keras.metrics.HitsMetric.submodules": true,
     "tfr.keras.metrics.HitsMetric.supports_masking": true,
@@ -2697,7 +2771,6 @@
     "tfr.keras.metrics.MRRMetric.__new__": true,
     "tfr.keras.metrics.MRRMetric.activity_regularizer": true,
     "tfr.keras.metrics.MRRMetric.add_loss": true,
-    "tfr.keras.metrics.MRRMetric.add_metric": true,
     "tfr.keras.metrics.MRRMetric.build": true,
     "tfr.keras.metrics.MRRMetric.build_from_config": true,
     "tfr.keras.metrics.MRRMetric.compute_dtype": true,
@@ -2713,6 +2786,7 @@
     "tfr.keras.metrics.MRRMetric.get_weights": true,
     "tfr.keras.metrics.MRRMetric.input": true,
     "tfr.keras.metrics.MRRMetric.input_spec": true,
+    "tfr.keras.metrics.MRRMetric.load_own_variables": true,
     "tfr.keras.metrics.MRRMetric.losses": true,
     "tfr.keras.metrics.MRRMetric.merge_state": true,
     "tfr.keras.metrics.MRRMetric.metrics": true,
@@ -2722,6 +2796,7 @@
     "tfr.keras.metrics.MRRMetric.output": true,
     "tfr.keras.metrics.MRRMetric.reset_state": true,
     "tfr.keras.metrics.MRRMetric.result": true,
+    "tfr.keras.metrics.MRRMetric.save_own_variables": true,
     "tfr.keras.metrics.MRRMetric.set_weights": true,
     "tfr.keras.metrics.MRRMetric.submodules": true,
     "tfr.keras.metrics.MRRMetric.supports_masking": true,
@@ -2743,7 +2818,6 @@
     "tfr.keras.metrics.MeanAveragePrecisionMetric.__new__": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.activity_regularizer": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.add_loss": true,
-    "tfr.keras.metrics.MeanAveragePrecisionMetric.add_metric": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.build": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.build_from_config": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.compute_dtype": true,
@@ -2759,6 +2833,7 @@
     "tfr.keras.metrics.MeanAveragePrecisionMetric.get_weights": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.input": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.input_spec": true,
+    "tfr.keras.metrics.MeanAveragePrecisionMetric.load_own_variables": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.losses": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.merge_state": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.metrics": true,
@@ -2768,6 +2843,7 @@
     "tfr.keras.metrics.MeanAveragePrecisionMetric.output": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.reset_state": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.result": true,
+    "tfr.keras.metrics.MeanAveragePrecisionMetric.save_own_variables": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.set_weights": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.submodules": true,
     "tfr.keras.metrics.MeanAveragePrecisionMetric.supports_masking": true,
@@ -2789,7 +2865,6 @@
     "tfr.keras.metrics.NDCGMetric.__new__": true,
     "tfr.keras.metrics.NDCGMetric.activity_regularizer": true,
     "tfr.keras.metrics.NDCGMetric.add_loss": true,
-    "tfr.keras.metrics.NDCGMetric.add_metric": true,
     "tfr.keras.metrics.NDCGMetric.build": true,
     "tfr.keras.metrics.NDCGMetric.build_from_config": true,
     "tfr.keras.metrics.NDCGMetric.compute_dtype": true,
@@ -2805,6 +2880,7 @@
     "tfr.keras.metrics.NDCGMetric.get_weights": true,
     "tfr.keras.metrics.NDCGMetric.input": true,
     "tfr.keras.metrics.NDCGMetric.input_spec": true,
+    "tfr.keras.metrics.NDCGMetric.load_own_variables": true,
     "tfr.keras.metrics.NDCGMetric.losses": true,
     "tfr.keras.metrics.NDCGMetric.merge_state": true,
     "tfr.keras.metrics.NDCGMetric.metrics": true,
@@ -2814,6 +2890,7 @@
     "tfr.keras.metrics.NDCGMetric.output": true,
     "tfr.keras.metrics.NDCGMetric.reset_state": true,
     "tfr.keras.metrics.NDCGMetric.result": true,
+    "tfr.keras.metrics.NDCGMetric.save_own_variables": true,
     "tfr.keras.metrics.NDCGMetric.set_weights": true,
     "tfr.keras.metrics.NDCGMetric.submodules": true,
     "tfr.keras.metrics.NDCGMetric.supports_masking": true,
@@ -2835,7 +2912,6 @@
     "tfr.keras.metrics.OPAMetric.__new__": true,
     "tfr.keras.metrics.OPAMetric.activity_regularizer": true,
     "tfr.keras.metrics.OPAMetric.add_loss": true,
-    "tfr.keras.metrics.OPAMetric.add_metric": true,
     "tfr.keras.metrics.OPAMetric.build": true,
     "tfr.keras.metrics.OPAMetric.build_from_config": true,
     "tfr.keras.metrics.OPAMetric.compute_dtype": true,
@@ -2851,6 +2927,7 @@
     "tfr.keras.metrics.OPAMetric.get_weights": true,
     "tfr.keras.metrics.OPAMetric.input": true,
     "tfr.keras.metrics.OPAMetric.input_spec": true,
+    "tfr.keras.metrics.OPAMetric.load_own_variables": true,
     "tfr.keras.metrics.OPAMetric.losses": true,
     "tfr.keras.metrics.OPAMetric.merge_state": true,
     "tfr.keras.metrics.OPAMetric.metrics": true,
@@ -2860,6 +2937,7 @@
     "tfr.keras.metrics.OPAMetric.output": true,
     "tfr.keras.metrics.OPAMetric.reset_state": true,
     "tfr.keras.metrics.OPAMetric.result": true,
+    "tfr.keras.metrics.OPAMetric.save_own_variables": true,
     "tfr.keras.metrics.OPAMetric.set_weights": true,
     "tfr.keras.metrics.OPAMetric.submodules": true,
     "tfr.keras.metrics.OPAMetric.supports_masking": true,
@@ -2881,7 +2959,6 @@
     "tfr.keras.metrics.PrecisionIAMetric.__new__": true,
     "tfr.keras.metrics.PrecisionIAMetric.activity_regularizer": true,
     "tfr.keras.metrics.PrecisionIAMetric.add_loss": true,
-    "tfr.keras.metrics.PrecisionIAMetric.add_metric": true,
     "tfr.keras.metrics.PrecisionIAMetric.build": true,
     "tfr.keras.metrics.PrecisionIAMetric.build_from_config": true,
     "tfr.keras.metrics.PrecisionIAMetric.compute_dtype": true,
@@ -2897,6 +2974,7 @@
     "tfr.keras.metrics.PrecisionIAMetric.get_weights": true,
     "tfr.keras.metrics.PrecisionIAMetric.input": true,
     "tfr.keras.metrics.PrecisionIAMetric.input_spec": true,
+    "tfr.keras.metrics.PrecisionIAMetric.load_own_variables": true,
     "tfr.keras.metrics.PrecisionIAMetric.losses": true,
     "tfr.keras.metrics.PrecisionIAMetric.merge_state": true,
     "tfr.keras.metrics.PrecisionIAMetric.metrics": true,
@@ -2906,6 +2984,7 @@
     "tfr.keras.metrics.PrecisionIAMetric.output": true,
     "tfr.keras.metrics.PrecisionIAMetric.reset_state": true,
     "tfr.keras.metrics.PrecisionIAMetric.result": true,
+    "tfr.keras.metrics.PrecisionIAMetric.save_own_variables": true,
     "tfr.keras.metrics.PrecisionIAMetric.set_weights": true,
     "tfr.keras.metrics.PrecisionIAMetric.submodules": true,
     "tfr.keras.metrics.PrecisionIAMetric.supports_masking": true,
@@ -2927,7 +3006,6 @@
     "tfr.keras.metrics.PrecisionMetric.__new__": true,
     "tfr.keras.metrics.PrecisionMetric.activity_regularizer": true,
     "tfr.keras.metrics.PrecisionMetric.add_loss": true,
-    "tfr.keras.metrics.PrecisionMetric.add_metric": true,
     "tfr.keras.metrics.PrecisionMetric.build": true,
     "tfr.keras.metrics.PrecisionMetric.build_from_config": true,
     "tfr.keras.metrics.PrecisionMetric.compute_dtype": true,
@@ -2943,6 +3021,7 @@
     "tfr.keras.metrics.PrecisionMetric.get_weights": true,
     "tfr.keras.metrics.PrecisionMetric.input": true,
     "tfr.keras.metrics.PrecisionMetric.input_spec": true,
+    "tfr.keras.metrics.PrecisionMetric.load_own_variables": true,
     "tfr.keras.metrics.PrecisionMetric.losses": true,
     "tfr.keras.metrics.PrecisionMetric.merge_state": true,
     "tfr.keras.metrics.PrecisionMetric.metrics": true,
@@ -2952,6 +3031,7 @@
     "tfr.keras.metrics.PrecisionMetric.output": true,
     "tfr.keras.metrics.PrecisionMetric.reset_state": true,
     "tfr.keras.metrics.PrecisionMetric.result": true,
+    "tfr.keras.metrics.PrecisionMetric.save_own_variables": true,
     "tfr.keras.metrics.PrecisionMetric.set_weights": true,
     "tfr.keras.metrics.PrecisionMetric.submodules": true,
     "tfr.keras.metrics.PrecisionMetric.supports_masking": true,
@@ -2992,7 +3072,6 @@
     "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.build": true,
     "tfr.keras.metrics.RecallMetric.build_from_config": true,
     "tfr.keras.metrics.RecallMetric.compute_dtype": true,
@@ -3008,6 +3087,7 @@
     "tfr.keras.metrics.RecallMetric.get_weights": true,
     "tfr.keras.metrics.RecallMetric.input": true,
     "tfr.keras.metrics.RecallMetric.input_spec": true,
+    "tfr.keras.metrics.RecallMetric.load_own_variables": true,
     "tfr.keras.metrics.RecallMetric.losses": true,
     "tfr.keras.metrics.RecallMetric.merge_state": true,
     "tfr.keras.metrics.RecallMetric.metrics": true,
@@ -3017,6 +3097,7 @@
     "tfr.keras.metrics.RecallMetric.output": true,
     "tfr.keras.metrics.RecallMetric.reset_state": true,
     "tfr.keras.metrics.RecallMetric.result": true,
+    "tfr.keras.metrics.RecallMetric.save_own_variables": true,
     "tfr.keras.metrics.RecallMetric.set_weights": true,
     "tfr.keras.metrics.RecallMetric.submodules": true,
     "tfr.keras.metrics.RecallMetric.supports_masking": true,
@@ -3422,11 +3503,13 @@
     "tfr.keras.utils.PositiveFunction": false,
     "tfr.keras.utils.RankDiscountFunction": false,
     "tfr.keras.utils.TensorLike": false,
+    "tfr.keras.utils.deserialize_keras_object": 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.keras.utils.serialize_keras_object": false,
     "tfr.keras.utils.symmetric_log1p": false,
     "tfr.utils": false,
     "tfr.utils.LossFunction": false,
diff --git a/docs/api_docs/python/tfr/_toc.yaml b/docs/api_docs/python/tfr/_toc.yaml
index 3505887..300540d 100644
--- a/docs/api_docs/python/tfr/_toc.yaml
+++ b/docs/api_docs/python/tfr/_toc.yaml
@@ -153,6 +153,10 @@ toc:
       path: /ranking/api_docs/python/tfr/keras/losses/SoftmaxLoss
     - title: UniqueSoftmaxLoss
       path: /ranking/api_docs/python/tfr/keras/losses/UniqueSoftmaxLoss
+    - title: YetiDCGLambdaWeight
+      path: /ranking/api_docs/python/tfr/keras/losses/YetiDCGLambdaWeight
+    - title: YetiLogisticLoss
+      path: /ranking/api_docs/python/tfr/keras/losses/YetiLogisticLoss
     - title: get
       path: /ranking/api_docs/python/tfr/keras/losses/get
   - title: metrics
@@ -273,6 +277,8 @@ toc:
       path: /ranking/api_docs/python/tfr/keras/utils
     - title: GainFunction
       path: /ranking/api_docs/python/tfr/keras/utils/GainFunction
+    - title: deserialize_keras_object
+      path: /ranking/api_docs/python/tfr/keras/utils/deserialize_keras_object
     - title: identity
       path: /ranking/api_docs/python/tfr/keras/utils/identity
     - title: inverse
@@ -283,6 +289,8 @@ toc:
       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: serialize_keras_object
+      path: /ranking/api_docs/python/tfr/keras/utils/serialize_keras_object
     - title: symmetric_log1p
       path: /ranking/api_docs/python/tfr/keras/utils/symmetric_log1p
 - title: tfr.utils
diff --git a/docs/api_docs/python/tfr/all_symbols.md b/docs/api_docs/python/tfr/all_symbols.md
index ae880dc..1195512 100644
--- a/docs/api_docs/python/tfr/all_symbols.md
+++ b/docs/api_docs/python/tfr/all_symbols.md
@@ -81,6 +81,8 @@
 *   <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/YetiDCGLambdaWeight.md"><code>tfr.keras.losses.YetiDCGLambdaWeight</code></a>
+*   <a href="../tfr/keras/losses/YetiLogisticLoss.md"><code>tfr.keras.losses.YetiLogisticLoss</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>
@@ -138,11 +140,13 @@
 *   <a href="../tfr/keras/utils/GainFunction.md"><code>tfr.keras.utils.PositiveFunction</code></a>
 *   <a href="../tfr/keras/utils/GainFunction.md"><code>tfr.keras.utils.RankDiscountFunction</code></a>
 *   <a href="../tfr/keras/model/TensorLike.md"><code>tfr.keras.utils.TensorLike</code></a>
+*   <a href="../tfr/keras/utils/deserialize_keras_object.md"><code>tfr.keras.utils.deserialize_keras_object</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/keras/utils/serialize_keras_object.md"><code>tfr.keras.utils.serialize_keras_object</code></a>
 *   <a href="../tfr/keras/utils/symmetric_log1p.md"><code>tfr.keras.utils.symmetric_log1p</code></a>
 *   <a href="../tfr/utils.md"><code>tfr.utils</code></a>
 *   <a href="../tfr/utils/LossFunction.md"><code>tfr.utils.LossFunction</code></a>
diff --git a/docs/api_docs/python/tfr/api_report.pb b/docs/api_docs/python/tfr/api_report.pb
index c8e1ebe4f3ab8143da4471da37a8a704a4bf2ca2..0fef1f13f4a71d7255d7a2cdef4ec20583341d27 100644
GIT binary patch
delta 4271
zcmb7HZBSEJ8s=j73b`SK#8FU87022x7(h_@3_>M<l34s4tw1p{2pU#EMN3JG?XDu!
z<zY733Wb(&uu?3&v~Cq+uoYWrCG^L3yIptK&a|^TTR(PMZ3kuCbM7VK8Z4ndGV`8u
z-uKP>oaZ_3IrrviiRFr9u~g=N?VbO;?W?eK-KuAlVKLg6*hp=Bq&8+<j5cBEvIK2>
zRLqLl*yx_cFZ#-iX>+~2EU62<yb97_iJr?je3K+H*-GcEVkMo0!*>pXlp0)@1#&M0
z>O%s(Tpki`C_GITC`!bIJ-tuDSkPV_j@y4{VOx@q(JDEf)&}9B)(5AxD_~4eK@l53
z#wrZGtA_o!4>Jqx>lYz4TS-^A9noJhC@+>U$1z;B7&~G=LuosUr?iUTGIQbfeDl8D
z`FjiZloUQ+X2S37v#kX|jE{F`8eUS&#F=Qh%e$g*uj$vlqxVKYQ!eL{mueVNvmrMe
zIxCCXwSG7jw+-9kOJ`I`qfr>Y5PvlUiE8c1T8M2YSlo}BiJ@iw7&w~4buDWb*X+11
zEVr(BA78)4+taMw>Fu$0&X}Ng*1hCg7KzCb_;;cKA0$<K*D+X<oZ(jTjql~WQ$dPQ
za#YGU;*tkX2{}e|5|rusB!?KO4$;X`)2n;J-E-kjrmYhSkAWjYh6^hV)&>c~VoPQS
z9;W-?WO_PGn^pYc!*sDq=+1{3>57W(IvR~C&XMtw7{;Cba^yej%Ll)?>MVZuHx|ho
z{pfglNT9P%M0f-PO&(t;w>wi#LDlGwU~G#VwOJc^rJvOY;KwX62BLH>HT!3A%{yXM
zppJF1vtZFJ!L{-6cdPZhmb0^ei}DlJ1S_T`{KPdO0uK$s{=Z=iL5V@n@BN)z8HNqw
zy2;*iVMc#%(a99{y5F_OB<MHpF^%xJ$DAfFlJU0akakan_g-O9zOD)v_sG%xj4xU@
zDe>ZZz4g97QIK`hZ2Yj^t!mt|vACY#M_u3Ex<L%!*Z_-ptLVzr8P(D_G;CThJpe9m
z&&8WIr#^wXS~+X3s4!JTZ8lYwWTa=JfBP)IHeDoX!8(j=&ydi07z#F`)}#^!^3r^y
z6si225y3kWV3`?=@)dzrhcCkhmQ7?tQu!>LFO(_0XjgwS{V6Hy<>jSQlOO)IYc|FP
z*btHXYIAW(MP*45EL%6CVW&z)|HC!evt>0p4zw|Lz9-&2qK3vI#hrr^BwfxyI46v^
zzMx0er-HrysQaGltinob4Q)0BGy~UKvq!9YN2B2KTx~as&hxhC&9@$8nXsltf0FQ?
zuW<8VIA^gd!IXShhb2rj8oFt;6vg&jur>2=^>7|;I)qv_^l!k4m)DL%Pi$5r_0=G>
z)d|pV)|HWNF3i;&a)&%Nz^`!sJ@kthW?})0@IPE*_CedPqUk(?<S(&ZJ;EQ!{wM}2
zN;4%?(A%xFT}VNZj#r1!o;-ve>2s00n<uobL5f@T0-@A?JfUq3GR}VVTQNFE;}`ti
zlOWHy#>rEn@3kM?$)gFTL^MGjEl?m&4Qf+Jlt5}z0_W^!aiLj>?4}3;)7&)P-ugcC
zhp{;jhoR?7ruHP9Tw?PhEzZ)d=7W~oGEv_k*uVT)Pg<TFxpqIg<{3B7XWKI@+WsPx
zP++|!#;>#UiKoWTGpfmBvBnz0wH&Wz#wm<@U7&E}q!dwa2ozF#5;XF<QDfw!?ERCE
zFsNu17$h;Ck?;&c*P_IO)=3o74AZ2L7U-mq+RigbZAxKjdme#zslA1qxBZfr8EFdA
zIwMj4vHPJrTO>#F=>@QK2#*>=$5&!`)|y!4bS^;FKqdxH?U&Rs(o~wsGg6+jfBYmz
zYM(#ubc8{D=~<FOixXG%Na@TF@i&omO7uK|%aT~m)_IANXKspx=1dUm@8k&gFBPcf
zOkh)&4_eNo3oVTf-7(X)MEr20i^q&^Va(m=suE+jbPu^?O5CRj*(p={2j;VrWJ(ix
zMk&aYN5gf!=aMSDN^bW`vHY@ube&@U`k>>-OZ9fIj0@_0M62}5paCf#3RecPJ;QQc
z88kBR+$7o3BvV3~=sYxoQ?}H0%9h$xwg*01M~EHzs2N|KH}h|V|Jfq2>w-Ijp=2YB
zR{~({3+BGRFeYa0i%-JIbO4syKE%8+^5tL)^$d|0YsAGxM5GTy!klqH8V^#z)bMy+
zl77bJ+D|!~ie?K7&E;Y|XrhHG5dN{Azj?^_;1$2e9K%c1EX>2Cxe;IIQC#!MX)&;d
zoh<Ah2rgLzy~2a-zAcz$oE}1MY+x=Ug7;{ODJi5FxvWnEyvLhn%WG;l23cN-yS^r0
z41-&lN)+~zxBTD}Zr@Ee9rI!PTsUeqRV<SF{b?*``~097-^CJg>kO%J_<9dhh4xoz
zDt4N{5{6V_V4`@>sA?SF%}4&-I9Pv-!m4|7sLd&Vn$OIzbfjBDmCSq;#QJjicSj^R
z_$QflM9KIto3M6t4*!<^rj$kfKLfa9qdP;$4aY93sw^q1h}v$dFp)3pvXTR){8H1N
z!is$J^E*vNl?v}rAM&ask5=vKrRPRjW|V=WkqsB3ek}&aG+H=H*>m@Tn4cy82b~Ir
A9RL6T

delta 3131
zcmb7GX;4#F6y~wS5b{D;qAU`lwzXnK3@U;m0YqbANJR&+jEb#Gr3z}b>VO!ERtGHz
zUu|3{sKeM=l|*mB1sAG=)he2zb*rd9oZ8N~;I!jtTebJT$4es9GEM&Ey>st5=eyr`
z&biOr;{9N^_dtbG{`|l%H8QoW{9Xwc7_N`j#|(}hIyh$NWW7FqSWJ9$T$Damuh;J%
zSS(W-k|I1j%!#jic)XK@%Gd~EmyPyT%9AClGMQOnm-U8!pbi(5ex%6H807cV;pagv
zYLk?HE`dn+pk?0Nb`<1kVY|oU_>dWxU#!Ag!~GRAXU&?iVD^&5^WL}Y@Z};xbt#bw
z-h=)$9v%imyyMvnb@T^bvlWBTcX${!MF*o&r$ugb2>A41*z~?cHnfd{VOA7mF=vr6
zTLtq7nXq6V7QoT$J6g3>BNlt*D<ZKdUe`5o*ncKA;O0kRG}ts8pQl1myrwGFncTFK
zM`3(VJcyN(Rk8O)0JXWoJ{*lZK&Fo}b|y#1Lmt`5ui?pPlUI>qBr=kt+(EH%W2oTX
zk>G>EgcNQCniIx4Sg843#a7A3W1>6-7VK6Yw+(m6kI9oe0a$K_ay+6Y_JF~t#H_I<
z%aH(TSU(}CD~xVb@NWo1^Vl3qxLkC_s8k0Vq#{)=CV!cyf^D*lCQ{?K@q0+|geV4L
z-`iz^@Zst12p^fYSP)*F9)`nd&TM@8Gf8k-0cpPI9}|MaARp+GB9Rv(+8Q#9&SFUO
zB`-|gUFdjEq+vy-`!w``Ia3MSM59Ot4MAUg5RvKc(4`_Zn3yyF5P^!ce23x*!OW8s
zX7QZocEhF;E*J~P@u+X-u`N3k`W6%XN`>dH(sWm+rYs#6<^*C~wvo$6MRuHog2<=7
z(J9PI6{@pC@JUVs`dEYvWYV<IzF1bWox8(&1UHXrpet42X_W~-mGy*qwXZuM`kF(r
zzBCYDw-~v(ux-$=0EjGg=vFm?$Ew3}A&38J`rnvFI!Y@=hOBE^vF(-!z<=$1PeT$O
z-Rpybbd{)*;ey~uj;;HY>%?1;%A=rc4yx054BjfoqjGnO>rB}Sth!74wY(>Kopyjk
zQ=5iGBJJy+ipcI?6Aai~GamWn^I_U4zAvkI;(`(dCLWKz5m1o0$`T1ABh|cFB2{>H
zm((Wo!i%yLfmB*PG*~MfcqLM)nHWC5E0HNCWsqOy@yK5*7yPb%VH8MaF_u+iqJ5^c
z#^c>fsmMRc{Io8|a+1#y<ZfBsc`k9J{JbrQRf}_zuEwwH>AD>${Xjif#fW{>ev*nz
z4G*_D&IhV8F;)L1^^pizSL2WBnmy9tna|DIhc390lf@(Agb<FD?^xj~hq}%w2Yv1e
zxVuJUHC|Zftnpx|>$Fm{NoZL7mvv%&->K`Y@7i5reG7|U{d;}WHL`E_I*DSd7oD^T
z8{{ac4`(!c6)Iw}s4fT#>%+Q&?uO;ld)9a@SB#N7QMA`VF0t>;5y!(+>m!(-tI%%k
zk2`Id7Ndsdp1$AeLMb_U;JgG)xKKtN3PS6*F6x^b-BW+5Q+4Y4N*A);u5Ah;A&tMj
zsyy46amY%!CaVr}oC`jO<ZMe^Wd^jvc~aY?bxNO4$C?KAX7?pO83_9^ALJjCx{wB;
z76*}%^voR6I0+r4LLM|;%pvL%y_{WXZ>D3oCdJW}_U25XBblRXN3!=RmyEO{V_c5t
zY&5pM>v9I_w+M-d{Zw!1oNPYnT4thkGyCr;`+px<OInJ!0z7Vsm#P^rDmBD>x|diV
z%7r|%_5d`W^~L=&jkxyx00Hm0>nuqij{i{Vf)mywqM1Sb@bmz(?tF*{uL{pI#5rki
zQ+F-mZaU9WAHapea3PmlfE^b`IzX;W=MjBN4jPS&x%jt8bbdJ{rA=LB^TuaZ%JoVf
ziI+mEY(~=57V%2S_%WTR>#oapCHmMUbmuycb5~3#X%U{?z=_=AP^7hMF#4*I%fn-O
z`>|lI7mRMJ5R~-3)*-@ONk0G*zLe9A0m!|v4<}kh1gsva^*n4hGzhy{?7}a-<eOWC
z$RxM@R6lgI@Rp-q+<;9}xFJ}Y7h<8KV_3r#u0Mx4K00*F^%E=yvpu=J-bhddcxI(Q
ZXM6KtVh_$7bZ|g8_a}F7Z3Oqs`yW|nSx5i?

diff --git a/docs/api_docs/python/tfr/data/parse_from_example_in_example.md b/docs/api_docs/python/tfr/data/parse_from_example_in_example.md
index ad0d5d7..3c5cd11 100644
--- a/docs/api_docs/python/tfr/data/parse_from_example_in_example.md
+++ b/docs/api_docs/python/tfr/data/parse_from_example_in_example.md
@@ -146,17 +146,14 @@ And the expected output is:
 
 ```python
 {
-    "unigrams":
-        SparseTensor(
-            indices=array([[0, 0, 0], [0, 1, 0], [0, 1, 1], [0, 1, 2],
-                           [1, 0, 0], [1, 1, 0], [1, 1, 1]]),
-            values=[
-                "tensorflow", "learning", "to", "rank", "gbdt", "neural",
-                "network"
-            ],
-            dense_shape=array([2, 2, 3])),
-    "utility": [[[0.], [1.]], [[0.], [1.]]],
-    "query_length": [[3], [2]],
+  "unigrams": SparseTensor(
+    indices=array([[0, 0, 0], [0, 1, 0], [0, 1, 1], [0, 1, 2], [1, 0, 0],
+      [1, 1, 0], [1, 1, 1]]),
+    values=["tensorflow", "learning", "to", "rank", "gbdt", "neural" ,
+      "network"],
+    dense_shape=array([2, 2, 3])),
+  "utility": [[[ 0.], [ 1.]], [[ 0.], [ 1.]]],
+  "query_length": [[3], [2]],
 }
 ```
 
diff --git a/docs/api_docs/python/tfr/data/parse_from_example_list.md b/docs/api_docs/python/tfr/data/parse_from_example_list.md
index 6d61a3e..79c6566 100644
--- a/docs/api_docs/python/tfr/data/parse_from_example_list.md
+++ b/docs/api_docs/python/tfr/data/parse_from_example_list.md
@@ -126,17 +126,14 @@ And the expected output is:
 
 ```python
 {
-    "unigrams":
-        SparseTensor(
-            indices=array([[0, 0, 0], [0, 1, 0], [0, 1, 1], [0, 1, 2],
-                           [1, 0, 0], [1, 1, 0], [1, 1, 1]]),
-            values=[
-                "tensorflow", "learning", "to", "rank", "gbdt", "neural",
-                "network"
-            ],
-            dense_shape=array([2, 2, 3])),
-    "utility": [[[0.], [1.]], [[0.], [1.]]],
-    "query_length": [[3], [2]],
+  "unigrams": SparseTensor(
+    indices=array([[0, 0, 0], [0, 1, 0], [0, 1, 1], [0, 1, 2], [1, 0, 0],
+      [1, 1, 0], [1, 1, 1]]),
+    values=["tensorflow", "learning", "to", "rank", "gbdt", "neural" ,
+      "network"],
+    dense_shape=array([2, 2, 3])),
+  "utility": [[[ 0.], [ 1.]], [[ 0.], [ 1.]]],
+  "query_length": [[3], [2]],
 }
 ```
 
diff --git a/docs/api_docs/python/tfr/data/parse_from_sequence_example.md b/docs/api_docs/python/tfr/data/parse_from_sequence_example.md
index 4460f39..bf1a3c8 100644
--- a/docs/api_docs/python/tfr/data/parse_from_sequence_example.md
+++ b/docs/api_docs/python/tfr/data/parse_from_sequence_example.md
@@ -119,14 +119,13 @@ And the expected output is:
 
 ```python
 {
-    "unigrams":
-        SparseTensor(
-            indices=array([[0, 0, 0], [0, 1, 0], [0, 1, 1], [0, 1, 2],
-                           [1, 0, 0], [1, 1, 0], [1, 1, 1]]),
-            values=["tensorflow", "learning", "to", "rank", "gbdt"],
-            dense_shape=array([2, 2, 3])),
-    "utility": [[[0.], [1.]], [[0.], [0.]]],
-    "query_length": [[3], [2]],
+  "unigrams": SparseTensor(
+    indices=array([[0, 0, 0], [0, 1, 0], [0, 1, 1], [0, 1, 2], [1, 0, 0], [1,
+    1, 0], [1, 1, 1]]),
+    values=["tensorflow", "learning", "to", "rank", "gbdt"],
+    dense_shape=array([2, 2, 3])),
+  "utility": [[[ 0.], [ 1.]], [[ 0.], [ 0.]]],
+  "query_length": [[3], [2]],
 }
 ```
 
diff --git a/docs/api_docs/python/tfr/extension/premade/TFRBertConfig.md b/docs/api_docs/python/tfr/extension/premade/TFRBertConfig.md
index 87c1a9c..fd482c7 100644
--- a/docs/api_docs/python/tfr/extension/premade/TFRBertConfig.md
+++ b/docs/api_docs/python/tfr/extension/premade/TFRBertConfig.md
@@ -19,6 +19,7 @@ description: The tf-ranking BERT task config.
 <meta itemprop="property" content="RESERVED_ATTR"/>
 <meta itemprop="property" content="SEQUENCE_TYPES"/>
 <meta itemprop="property" content="aggregated_metrics"/>
+<meta itemprop="property" content="allow_image_summary"/>
 <meta itemprop="property" content="default_params"/>
 <meta itemprop="property" content="differential_privacy_config"/>
 <meta itemprop="property" content="init_checkpoint"/>
@@ -68,6 +69,7 @@ Inherits From:
     validation_data: <a href="../../../tfr/extension/task/RankingDataConfig.md"><code>tfr.extension.task.RankingDataConfig</code></a> = None,
     name: Optional[str] = None,
     differential_privacy_config: Optional[dp_configs.DifferentialPrivacyConfig] = None,
+    allow_image_summary: bool = False,
     loss: str = &#x27;softmax_loss&#x27;,
     loss_reduction: str = tf.keras.losses.Reduction.NONE,
     aggregated_metrics: bool = False,
@@ -78,7 +80,6 @@ Inherits From:
 <!-- 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>
@@ -105,7 +106,6 @@ operators, including {'==', '!=',  '<', '<=', '>', '>='}.
 </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>
@@ -171,6 +171,13 @@ Dataclass field
 </td>
 </tr><tr>
 <td>
+`allow_image_summary`<a id="allow_image_summary"></a>
+</td>
+<td>
+Dataclass field
+</td>
+</tr><tr>
+<td>
 `loss`<a id="loss"></a>
 </td>
 <td>
@@ -424,6 +431,13 @@ aggregated_metrics<a id="aggregated_metrics"></a>
 </td>
 </tr><tr>
 <td>
+allow_image_summary<a id="allow_image_summary"></a>
+</td>
+<td>
+`False`
+</td>
+</tr><tr>
+<td>
 default_params<a id="default_params"></a>
 </td>
 <td>
diff --git a/docs/api_docs/python/tfr/extension/premade/TFRBertDataConfig.md b/docs/api_docs/python/tfr/extension/premade/TFRBertDataConfig.md
index e50689c..6671bba 100644
--- a/docs/api_docs/python/tfr/extension/premade/TFRBertDataConfig.md
+++ b/docs/api_docs/python/tfr/extension/premade/TFRBertDataConfig.md
@@ -19,6 +19,7 @@ description: Data config for TFR-BERT task.
 <meta itemprop="property" content="RESERVED_ATTR"/>
 <meta itemprop="property" content="SEQUENCE_TYPES"/>
 <meta itemprop="property" content="apply_tf_data_service_before_batching"/>
+<meta itemprop="property" content="autotune_algorithm"/>
 <meta itemprop="property" content="block_length"/>
 <meta itemprop="property" content="cache"/>
 <meta itemprop="property" content="convert_labels_to_binary"/>
@@ -107,6 +108,7 @@ Inherits From:
     trainer_id: Optional[str] = None,
     seed: Optional[int] = None,
     prefetch_buffer_size: Optional[int] = None,
+    autotune_algorithm: Optional[str] = None,
     data_format: str = tfr_data.ELWC,
     dataset_fn: str = &#x27;tfrecord&#x27;,
     list_size: Optional[int] = None,
@@ -122,7 +124,6 @@ Inherits From:
 <!-- 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>
@@ -149,7 +150,6 @@ operators, including {'==', '!=',  '<', '<=', '>', '>='}.
 </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>
@@ -334,6 +334,13 @@ Dataclass field
 </td>
 </tr><tr>
 <td>
+`autotune_algorithm`<a id="autotune_algorithm"></a>
+</td>
+<td>
+Dataclass field
+</td>
+</tr><tr>
+<td>
 `data_format`<a id="data_format"></a>
 </td>
 <td>
@@ -622,6 +629,13 @@ apply_tf_data_service_before_batching<a id="apply_tf_data_service_before_batchin
 </td>
 </tr><tr>
 <td>
+autotune_algorithm<a id="autotune_algorithm"></a>
+</td>
+<td>
+`None`
+</td>
+</tr><tr>
+<td>
 block_length<a id="block_length"></a>
 </td>
 <td>
diff --git a/docs/api_docs/python/tfr/extension/premade/TFRBertModelConfig.md b/docs/api_docs/python/tfr/extension/premade/TFRBertModelConfig.md
index ba08b3d..d584407 100644
--- a/docs/api_docs/python/tfr/extension/premade/TFRBertModelConfig.md
+++ b/docs/api_docs/python/tfr/extension/premade/TFRBertModelConfig.md
@@ -59,7 +59,6 @@ A TFR-BERT model configuration.
 <!-- 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>
@@ -86,7 +85,6 @@ operators, including {'==', '!=',  '<', '<=', '>', '>='}.
 </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>
diff --git a/docs/api_docs/python/tfr/extension/premade/TensorDict.md b/docs/api_docs/python/tfr/extension/premade/TensorDict.md
index 495b2f4..aa049ef 100644
--- a/docs/api_docs/python/tfr/extension/premade/TensorDict.md
+++ b/docs/api_docs/python/tfr/extension/premade/TensorDict.md
@@ -6,7 +6,6 @@
 # tfr.extension.premade.TensorDict
 
 <!-- Insert buttons and diff -->
-
 This symbol is a **type alias**.
 
 #### Source:
diff --git a/docs/api_docs/python/tfr/extension/task/FeatureSpec.md b/docs/api_docs/python/tfr/extension/task/FeatureSpec.md
index 3cf59ae..3574194 100644
--- a/docs/api_docs/python/tfr/extension/task/FeatureSpec.md
+++ b/docs/api_docs/python/tfr/extension/task/FeatureSpec.md
@@ -6,7 +6,6 @@
 # tfr.extension.task.FeatureSpec
 
 <!-- Insert buttons and diff -->
-
 This symbol is a **type alias**.
 
 #### Source:
diff --git a/docs/api_docs/python/tfr/extension/task/RankingDataConfig.md b/docs/api_docs/python/tfr/extension/task/RankingDataConfig.md
index d3ce606..3a84f30 100644
--- a/docs/api_docs/python/tfr/extension/task/RankingDataConfig.md
+++ b/docs/api_docs/python/tfr/extension/task/RankingDataConfig.md
@@ -19,6 +19,7 @@ description: Data set config.
 <meta itemprop="property" content="RESERVED_ATTR"/>
 <meta itemprop="property" content="SEQUENCE_TYPES"/>
 <meta itemprop="property" content="apply_tf_data_service_before_batching"/>
+<meta itemprop="property" content="autotune_algorithm"/>
 <meta itemprop="property" content="block_length"/>
 <meta itemprop="property" content="cache"/>
 <meta itemprop="property" content="convert_labels_to_binary"/>
@@ -93,6 +94,7 @@ Data set config.
     trainer_id: Optional[str] = None,
     seed: Optional[int] = None,
     prefetch_buffer_size: Optional[int] = None,
+    autotune_algorithm: Optional[str] = None,
     data_format: str = tfr_data.ELWC,
     dataset_fn: str = &#x27;tfrecord&#x27;,
     list_size: Optional[int] = None,
@@ -105,7 +107,6 @@ Data set config.
 <!-- 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>
@@ -132,7 +133,6 @@ operators, including {'==', '!=',  '<', '<=', '>', '>='}.
 </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>
@@ -317,6 +317,13 @@ Dataclass field
 </td>
 </tr><tr>
 <td>
+`autotune_algorithm`<a id="autotune_algorithm"></a>
+</td>
+<td>
+Dataclass field
+</td>
+</tr><tr>
+<td>
 `data_format`<a id="data_format"></a>
 </td>
 <td>
@@ -584,6 +591,13 @@ apply_tf_data_service_before_batching<a id="apply_tf_data_service_before_batchin
 </td>
 </tr><tr>
 <td>
+autotune_algorithm<a id="autotune_algorithm"></a>
+</td>
+<td>
+`None`
+</td>
+</tr><tr>
+<td>
 block_length<a id="block_length"></a>
 </td>
 <td>
diff --git a/docs/api_docs/python/tfr/extension/task/RankingTaskConfig.md b/docs/api_docs/python/tfr/extension/task/RankingTaskConfig.md
index 3fd85c5..3b52a78 100644
--- a/docs/api_docs/python/tfr/extension/task/RankingTaskConfig.md
+++ b/docs/api_docs/python/tfr/extension/task/RankingTaskConfig.md
@@ -19,6 +19,7 @@ description: The TF-Ranking task config.
 <meta itemprop="property" content="RESERVED_ATTR"/>
 <meta itemprop="property" content="SEQUENCE_TYPES"/>
 <meta itemprop="property" content="aggregated_metrics"/>
+<meta itemprop="property" content="allow_image_summary"/>
 <meta itemprop="property" content="default_params"/>
 <meta itemprop="property" content="differential_privacy_config"/>
 <meta itemprop="property" content="init_checkpoint"/>
@@ -57,6 +58,7 @@ The TF-Ranking task config.
     validation_data: <a href="../../../tfr/extension/task/RankingDataConfig.md"><code>tfr.extension.task.RankingDataConfig</code></a> = None,
     name: Optional[str] = None,
     differential_privacy_config: Optional[dp_configs.DifferentialPrivacyConfig] = None,
+    allow_image_summary: bool = False,
     loss: str = &#x27;softmax_loss&#x27;,
     loss_reduction: str = tf.keras.losses.Reduction.NONE,
     aggregated_metrics: bool = False,
@@ -67,7 +69,6 @@ The TF-Ranking task config.
 <!-- 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>
@@ -94,7 +95,6 @@ operators, including {'==', '!=',  '<', '<=', '>', '>='}.
 </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>
@@ -160,6 +160,13 @@ Dataclass field
 </td>
 </tr><tr>
 <td>
+`allow_image_summary`<a id="allow_image_summary"></a>
+</td>
+<td>
+Dataclass field
+</td>
+</tr><tr>
+<td>
 `loss`<a id="loss"></a>
 </td>
 <td>
@@ -413,6 +420,13 @@ aggregated_metrics<a id="aggregated_metrics"></a>
 </td>
 </tr><tr>
 <td>
+allow_image_summary<a id="allow_image_summary"></a>
+</td>
+<td>
+`False`
+</td>
+</tr><tr>
+<td>
 default_params<a id="default_params"></a>
 </td>
 <td>
diff --git a/docs/api_docs/python/tfr/keras/layers/Bilinear.md b/docs/api_docs/python/tfr/keras/layers/Bilinear.md
index 08e5288..94e45f1 100644
--- a/docs/api_docs/python/tfr/keras/layers/Bilinear.md
+++ b/docs/api_docs/python/tfr/keras/layers/Bilinear.md
@@ -7,7 +7,6 @@ description: A Keras Layer makes bilinear interaction of two vectors.
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,6 +16,8 @@ description: A Keras Layer makes bilinear interaction of two vectors.
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="with_name_scope"/>
 </div>
@@ -187,20 +188,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -347,55 +336,33 @@ Used for backwards compatibility only.
 </tr>
 </table>
 
-<h3 id="add_metric"><code>add_metric</code></h3>
+<h3 id="build"><code>build</code></h3>
+
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L821-L847">View
+source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
-<code>add_metric(
-    value, name=None, **kwargs
+<code>build(
+    input_shape: tf.TensorShape
 )
 </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')
+See tf.keras.layers.Layer.
 
-  def call(self, inputs):
-    self.add_metric(self.mean(inputs))
-    self.add_metric(tf.reduce_sum(inputs), name='metric_2')
-    return inputs
-```
+<h3 id="build_from_config"><code>build_from_config</code></h3>
 
-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()`.
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>build_from_config(
+    config
+)
+</code></pre>
 
-```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')
-```
+Builds the layer's states with the supplied config dict.
 
-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')
-```
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
 
 <!-- Tabular view -->
  <table class="responsive fixed orange">
@@ -404,53 +371,14 @@ model.add_metric(tf.keras.metrics.Mean()(x), name='metric_1')
 
 <tr>
 <td>
-`value`
-</td>
-<td>
-Metric tensor.
-</td>
-</tr><tr>
-<td>
-`name`
-</td>
-<td>
-String metric name.
-</td>
-</tr><tr>
-<td>
-`**kwargs`
+`config`
 </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`.
+Dict containing the input shape associated with this layer.
 </td>
 </tr>
 </table>
 
-<h3 id="build"><code>build</code></h3>
-
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L821-L847">View
-source</a>
-
-<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
-<code>build(
-    input_shape: tf.TensorShape
-)
-</code></pre>
-
-See tf.keras.layers.Layer.
-
-<h3 id="build_from_config"><code>build_from_config</code></h3>
-
-<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
-<code>build_from_config(
-    config
-)
-</code></pre>
-
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -594,6 +522,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L870-L877">View
@@ -655,6 +607,64 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/layers/ConcatFeatures.md b/docs/api_docs/python/tfr/keras/layers/ConcatFeatures.md
index 13f8a58..21ab943 100644
--- a/docs/api_docs/python/tfr/keras/layers/ConcatFeatures.md
+++ b/docs/api_docs/python/tfr/keras/layers/ConcatFeatures.md
@@ -8,7 +8,6 @@ manner.
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -18,6 +17,8 @@ manner.
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="with_name_scope"/>
 </div>
@@ -213,20 +214,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -373,89 +362,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -498,6 +404,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -674,6 +603,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L356-L361">View
@@ -758,6 +711,64 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/layers/DocumentInteractionAttention.md b/docs/api_docs/python/tfr/keras/layers/DocumentInteractionAttention.md
index 10ea7d4..818fa15 100644
--- a/docs/api_docs/python/tfr/keras/layers/DocumentInteractionAttention.md
+++ b/docs/api_docs/python/tfr/keras/layers/DocumentInteractionAttention.md
@@ -7,7 +7,6 @@ description: Cross Document Interaction Attention layer.
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,6 +16,8 @@ description: Cross Document Interaction Attention layer.
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="with_name_scope"/>
 </div>
@@ -255,20 +256,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -415,89 +404,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L474-L526">View
@@ -539,6 +445,29 @@ correspond to `inputs` argument of call.
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -715,6 +644,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L564-L573">View
@@ -799,6 +752,64 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/layers/FlattenList.md b/docs/api_docs/python/tfr/keras/layers/FlattenList.md
index 50f9ec2..b4c39e5 100644
--- a/docs/api_docs/python/tfr/keras/layers/FlattenList.md
+++ b/docs/api_docs/python/tfr/keras/layers/FlattenList.md
@@ -7,7 +7,6 @@ description: Layer to flatten the example list.
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,6 +16,8 @@ description: Layer to flatten the example list.
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="with_name_scope"/>
 </div>
@@ -204,20 +205,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -364,89 +353,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -489,6 +395,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -665,6 +594,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L163-L168">View
@@ -749,6 +702,64 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/layers/GAMLayer.md b/docs/api_docs/python/tfr/keras/layers/GAMLayer.md
index ea947d4..2aee7ba 100644
--- a/docs/api_docs/python/tfr/keras/layers/GAMLayer.md
+++ b/docs/api_docs/python/tfr/keras/layers/GAMLayer.md
@@ -7,7 +7,6 @@ description: Defines a generalized additive model (GAM) layer.
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,6 +16,8 @@ description: Defines a generalized additive model (GAM) layer.
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="with_name_scope"/>
 </div>
@@ -298,20 +299,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -458,89 +447,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -583,6 +489,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -759,6 +688,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L777-L789">View
@@ -843,6 +796,64 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/layers/RestoreList.md b/docs/api_docs/python/tfr/keras/layers/RestoreList.md
index 1ae3e73..150c338 100644
--- a/docs/api_docs/python/tfr/keras/layers/RestoreList.md
+++ b/docs/api_docs/python/tfr/keras/layers/RestoreList.md
@@ -7,7 +7,6 @@ description: Output layer to restore listwise output shape.
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,6 +16,8 @@ description: Output layer to restore listwise output shape.
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="with_name_scope"/>
 </div>
@@ -181,20 +182,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -341,89 +330,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -466,6 +372,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -642,6 +571,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/layers.py#L253-L258">View
@@ -726,6 +679,64 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/layers/SelfAttentionMask.md b/docs/api_docs/python/tfr/keras/layers/SelfAttentionMask.md
index cc4d627..675374f 100644
--- a/docs/api_docs/python/tfr/keras/layers/SelfAttentionMask.md
+++ b/docs/api_docs/python/tfr/keras/layers/SelfAttentionMask.md
@@ -7,7 +7,6 @@ description: Create 3D attention mask from a 2D tensor mask.
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,6 +16,8 @@ description: Create 3D attention mask from a 2D tensor mask.
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="with_name_scope"/>
 </div>
@@ -162,20 +163,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -322,89 +311,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -447,6 +353,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -623,6 +552,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -704,6 +657,64 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses.md b/docs/api_docs/python/tfr/keras/losses.md
index 5143f3c..57b64fb 100644
--- a/docs/api_docs/python/tfr/keras/losses.md
+++ b/docs/api_docs/python/tfr/keras/losses.md
@@ -91,6 +91,12 @@ cross-entropy loss between `y_true` and `y_pred`.
 [`class UniqueSoftmaxLoss`](../../tfr/keras/losses/UniqueSoftmaxLoss.md):
 Computes unique softmax cross-entropy loss between `y_true` and `y_pred`.
 
+[`class YetiDCGLambdaWeight`](../../tfr/keras/losses/YetiDCGLambdaWeight.md):
+Keras serializable class for YetiDCGLambdaWeight.
+
+[`class YetiLogisticLoss`](../../tfr/keras/losses/YetiLogisticLoss.md): Computes
+Yeti logistic loss between `y_true` and `y_pred`.
+
 ## Functions
 
 [`get(...)`](../../tfr/keras/losses/get.md): Factory method to get a ranking
diff --git a/docs/api_docs/python/tfr/keras/losses/ApproxMRRLoss.md b/docs/api_docs/python/tfr/keras/losses/ApproxMRRLoss.md
index 309f213..1460050 100644
--- a/docs/api_docs/python/tfr/keras/losses/ApproxMRRLoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/ApproxMRRLoss.md
@@ -15,7 +15,7 @@ description: Computes approximate MRR loss between y_true and y_pred.
 
 <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#L819-L887">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L965-L1033">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -113,13 +113,14 @@ Retrieval Measures, Qin et al, 2008][qin2008]
 </td>
 <td>
 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 under a
-`tf.distribute.Strategy`, except via `Model.compile()` and
-`Model.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.
+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 under a `tf.distribute.Strategy`, except via
+`Model.compile()` and `Model.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>
@@ -135,7 +136,7 @@ Optional name for the instance.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L585-L592">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L728-L738">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -176,7 +177,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L575-L583">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L718-L726">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -187,7 +188,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L221-L229">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L248-L256">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/ApproxNDCGLoss.md b/docs/api_docs/python/tfr/keras/losses/ApproxNDCGLoss.md
index e431c60..9a0c16d 100644
--- a/docs/api_docs/python/tfr/keras/losses/ApproxNDCGLoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/ApproxNDCGLoss.md
@@ -15,7 +15,7 @@ description: Computes approximate NDCG loss between y_true and y_pred.
 
 <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#L890-L963">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1036-L1109">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -117,13 +117,14 @@ $$
 </td>
 <td>
 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 under a
-`tf.distribute.Strategy`, except via `Model.compile()` and
-`Model.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.
+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 under a `tf.distribute.Strategy`, except via
+`Model.compile()` and `Model.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>
@@ -139,7 +140,7 @@ Optional name for the instance.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L585-L592">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L728-L738">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -180,7 +181,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L575-L583">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L718-L726">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -191,7 +192,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L221-L229">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L248-L256">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/ClickEMLoss.md b/docs/api_docs/python/tfr/keras/losses/ClickEMLoss.md
index 3659d92..4094e38 100644
--- a/docs/api_docs/python/tfr/keras/losses/ClickEMLoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/ClickEMLoss.md
@@ -15,7 +15,7 @@ description: Computes click EM loss between y_true and y_pred.
 
 <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#L1070-L1144">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1216-L1290">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -173,7 +173,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1138-L1144">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1284-L1290">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -184,7 +184,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L221-L229">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L248-L256">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/CoupledRankDistilLoss.md b/docs/api_docs/python/tfr/keras/losses/CoupledRankDistilLoss.md
index 38072f9..c1a6fb3 100644
--- a/docs/api_docs/python/tfr/keras/losses/CoupledRankDistilLoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/CoupledRankDistilLoss.md
@@ -15,7 +15,7 @@ description: Computes the Rank Distil loss between y_true and y_pred.
 
 <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#L1385-L1474">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1531-L1620">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -86,7 +86,6 @@ The Coupled-RankDistil loss is defined as: $$ \mathcal{L}(y, s) = -\sum_{\pi}
 [reddi2021]: https://research.google/pubs/pub50695/
 
 <!-- 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>
@@ -156,7 +155,6 @@ all entries in the list.
 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>
@@ -172,7 +170,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>
@@ -186,7 +183,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1467-L1474">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1613-L1620">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -197,7 +194,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L221-L229">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L248-L256">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/DCGLambdaWeight.md b/docs/api_docs/python/tfr/keras/losses/DCGLambdaWeight.md
index de99c5a..f66eadf 100644
--- a/docs/api_docs/python/tfr/keras/losses/DCGLambdaWeight.md
+++ b/docs/api_docs/python/tfr/keras/losses/DCGLambdaWeight.md
@@ -15,7 +15,7 @@ description: Keras serializable class for DCG.
 
 <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#L107-L130">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L109-L132">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -85,7 +85,7 @@ LambdaMART.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L123-L130">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L125-L132">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/GumbelApproxNDCGLoss.md b/docs/api_docs/python/tfr/keras/losses/GumbelApproxNDCGLoss.md
index a42c7bd..c9310f2 100644
--- a/docs/api_docs/python/tfr/keras/losses/GumbelApproxNDCGLoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/GumbelApproxNDCGLoss.md
@@ -16,7 +16,7 @@ y_pred.
 
 <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#L966-L1067">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1112-L1213">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -125,13 +125,14 @@ $$
 </td>
 <td>
 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 under a
-`tf.distribute.Strategy`, except via `Model.compile()` and
-`Model.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.
+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 under a `tf.distribute.Strategy`, except via
+`Model.compile()` and `Model.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>
@@ -147,7 +148,7 @@ Optional name for the instance.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L585-L592">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L728-L738">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -188,7 +189,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1049-L1056">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1195-L1202">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -199,7 +200,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1058-L1067">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1204-L1213">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/LabelDiffLambdaWeight.md b/docs/api_docs/python/tfr/keras/losses/LabelDiffLambdaWeight.md
index 62ac7d5..23716d7 100644
--- a/docs/api_docs/python/tfr/keras/losses/LabelDiffLambdaWeight.md
+++ b/docs/api_docs/python/tfr/keras/losses/LabelDiffLambdaWeight.md
@@ -15,7 +15,7 @@ description: Keras serializable class for LabelDiffLambdaWeight.
 
 <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#L96-L104">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L98-L106">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -36,7 +36,7 @@ Keras serializable class for LabelDiffLambdaWeight.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L103-L104">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L105-L106">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/ListMLELambdaWeight.md b/docs/api_docs/python/tfr/keras/losses/ListMLELambdaWeight.md
index f3a7f05..aea61de 100644
--- a/docs/api_docs/python/tfr/keras/losses/ListMLELambdaWeight.md
+++ b/docs/api_docs/python/tfr/keras/losses/ListMLELambdaWeight.md
@@ -15,7 +15,7 @@ description: LambdaWeight for ListMLE cost function.
 
 <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#L190-L201">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L217-L228">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -52,7 +52,7 @@ LambdaWeight for ListMLE cost function.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L198-L201">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L225-L228">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -61,7 +61,7 @@ source</a>
 
 <h3 id="individual_weights"><code>individual_weights</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L447-L453">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L460-L466">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -74,7 +74,7 @@ See `_LambdaWeight`.
 
 <h3 id="pair_weights"><code>pair_weights</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L443-L445">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L456-L458">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/ListMLELoss.md b/docs/api_docs/python/tfr/keras/losses/ListMLELoss.md
index bba6532..eda113e 100644
--- a/docs/api_docs/python/tfr/keras/losses/ListMLELoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/ListMLELoss.md
@@ -15,7 +15,7 @@ description: Computes ListMLE loss between y_true and y_pred.
 
 <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#L734-L816">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L880-L962">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -153,7 +153,7 @@ False, this loss will accept dense tensors.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L585-L592">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L728-L738">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -194,7 +194,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L575-L583">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L718-L726">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -205,7 +205,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L221-L229">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L248-L256">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/MeanSquaredLoss.md b/docs/api_docs/python/tfr/keras/losses/MeanSquaredLoss.md
index d38cb25..a3a9a99 100644
--- a/docs/api_docs/python/tfr/keras/losses/MeanSquaredLoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/MeanSquaredLoss.md
@@ -15,7 +15,7 @@ description: Computes mean squared loss between y_true and y_pred.
 
 <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#L1275-L1326">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1421-L1472">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -142,7 +142,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L239-L242">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L266-L269">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -153,7 +153,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L221-L229">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L248-L256">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/MixtureEMLoss.md b/docs/api_docs/python/tfr/keras/losses/MixtureEMLoss.md
index 6a77a62..55891f9 100644
--- a/docs/api_docs/python/tfr/keras/losses/MixtureEMLoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/MixtureEMLoss.md
@@ -15,7 +15,7 @@ description: Computes mixture EM loss between y_true and y_pred.
 
 <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#L1147-L1222">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1293-L1368">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -145,7 +145,7 @@ False, this loss will accept dense tensors.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L585-L592">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L728-L738">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -186,7 +186,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1217-L1222">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1363-L1368">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -197,7 +197,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L221-L229">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L248-L256">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/NDCGLambdaWeight.md b/docs/api_docs/python/tfr/keras/losses/NDCGLambdaWeight.md
index 7f9b7f3..83b096b 100644
--- a/docs/api_docs/python/tfr/keras/losses/NDCGLambdaWeight.md
+++ b/docs/api_docs/python/tfr/keras/losses/NDCGLambdaWeight.md
@@ -15,7 +15,7 @@ description: Keras serializable class for NDCG.
 
 <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#L154-L169">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L181-L196">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -86,7 +86,7 @@ LambdaMART.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L123-L130">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L125-L132">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/NDCGLambdaWeightV2.md b/docs/api_docs/python/tfr/keras/losses/NDCGLambdaWeightV2.md
index 6a499ba..c75e7ac 100644
--- a/docs/api_docs/python/tfr/keras/losses/NDCGLambdaWeightV2.md
+++ b/docs/api_docs/python/tfr/keras/losses/NDCGLambdaWeightV2.md
@@ -15,7 +15,7 @@ description: Keras serializable class for NDCG LambdaWeight V2 for topn.
 
 <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#L133-L151">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L135-L153">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -75,7 +75,7 @@ Keras serializable class for NDCG LambdaWeight V2 for topn.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L146-L151">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L148-L153">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/OrdinalLoss.md b/docs/api_docs/python/tfr/keras/losses/OrdinalLoss.md
index d225493..8e443c9 100644
--- a/docs/api_docs/python/tfr/keras/losses/OrdinalLoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/OrdinalLoss.md
@@ -15,7 +15,7 @@ description: Computes the Ordinal loss between y_true and y_pred.
 
 <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#L1329-L1382">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1475-L1528">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -85,13 +85,14 @@ $$
 </td>
 <td>
 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 under a
-`tf.distribute.Strategy`, except via `Model.compile()` and
-`Model.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.
+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 under a `tf.distribute.Strategy`, except via
+`Model.compile()` and `Model.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>
@@ -145,7 +146,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L239-L242">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L266-L269">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -156,7 +157,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L221-L229">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L248-L256">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/PairwiseHingeLoss.md b/docs/api_docs/python/tfr/keras/losses/PairwiseHingeLoss.md
index 2e168c1..9b384ff 100644
--- a/docs/api_docs/python/tfr/keras/losses/PairwiseHingeLoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/PairwiseHingeLoss.md
@@ -15,7 +15,7 @@ description: Computes pairwise hinge loss between y_true and y_pred.
 
 <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#L291-L355">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L321-L385">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -126,7 +126,7 @@ False, this loss will accept dense tensors.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L269-L276">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L296-L306">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -167,7 +167,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L259-L267">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L286-L294">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -178,7 +178,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L221-L229">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L248-L256">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/PairwiseLogisticLoss.md b/docs/api_docs/python/tfr/keras/losses/PairwiseLogisticLoss.md
index a640d3f..0416c0a 100644
--- a/docs/api_docs/python/tfr/keras/losses/PairwiseLogisticLoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/PairwiseLogisticLoss.md
@@ -15,7 +15,7 @@ description: Computes pairwise logistic loss between y_true and y_pred.
 
 <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#L358-L422">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L388-L452">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -126,7 +126,7 @@ False, this loss will accept dense tensors.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L269-L276">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L296-L306">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -167,7 +167,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L259-L267">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L286-L294">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -178,7 +178,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L221-L229">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L248-L256">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/PairwiseMSELoss.md b/docs/api_docs/python/tfr/keras/losses/PairwiseMSELoss.md
index 9eb9c33..9744ee4 100644
--- a/docs/api_docs/python/tfr/keras/losses/PairwiseMSELoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/PairwiseMSELoss.md
@@ -16,7 +16,7 @@ y_pred.
 
 <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#L493-L558">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L523-L588">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -128,7 +128,7 @@ False, this loss will accept dense tensors.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L269-L276">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L296-L306">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -169,7 +169,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L259-L267">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L286-L294">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -180,7 +180,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L221-L229">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L248-L256">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/PairwiseSoftZeroOneLoss.md b/docs/api_docs/python/tfr/keras/losses/PairwiseSoftZeroOneLoss.md
index f0af422..d6822a2 100644
--- a/docs/api_docs/python/tfr/keras/losses/PairwiseSoftZeroOneLoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/PairwiseSoftZeroOneLoss.md
@@ -15,7 +15,7 @@ description: Computes pairwise soft zero-one loss between y_true and y_pred.
 
 <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#L425-L490">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L455-L520">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -127,7 +127,7 @@ False, this loss will accept dense tensors.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L269-L276">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L296-L306">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -168,7 +168,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L259-L267">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L286-L294">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -179,7 +179,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L221-L229">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L248-L256">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/PrecisionLambdaWeight.md b/docs/api_docs/python/tfr/keras/losses/PrecisionLambdaWeight.md
index 167bed0..16361ca 100644
--- a/docs/api_docs/python/tfr/keras/losses/PrecisionLambdaWeight.md
+++ b/docs/api_docs/python/tfr/keras/losses/PrecisionLambdaWeight.md
@@ -15,7 +15,7 @@ description: Keras serializable class for Precision.
 
 <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#L172-L187">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L199-L214">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -61,7 +61,7 @@ for positive examples. The rest are negative examples.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L183-L187">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L210-L214">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -118,7 +118,7 @@ A `Tensor` that can weight individual examples.
 
 <h3 id="pair_weights"><code>pair_weights</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L399-L427">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L412-L440">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/RankingLossKey.md b/docs/api_docs/python/tfr/keras/losses/RankingLossKey.md
index 9ffc302..8da9d24 100644
--- a/docs/api_docs/python/tfr/keras/losses/RankingLossKey.md
+++ b/docs/api_docs/python/tfr/keras/losses/RankingLossKey.md
@@ -18,6 +18,7 @@ description: Ranking loss key strings.
 <meta itemprop="property" content="SIGMOID_CROSS_ENTROPY_LOSS"/>
 <meta itemprop="property" content="SOFTMAX_LOSS"/>
 <meta itemprop="property" content="UNIQUE_SOFTMAX_LOSS"/>
+<meta itemprop="property" content="YETI_LOGISTIC_LOSS"/>
 </div>
 
 # tfr.keras.losses.RankingLossKey
@@ -26,7 +27,7 @@ description: Ranking loss key strings.
 
 <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#L11-L32">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L11-L33">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -41,7 +42,7 @@ Ranking loss key strings.
 
 <h3 id="all_keys"><code>all_keys</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L30-L32">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L31-L33">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -152,5 +153,12 @@ UNIQUE_SOFTMAX_LOSS<a id="UNIQUE_SOFTMAX_LOSS"></a>
 <td>
 `'unique_softmax_loss'`
 </td>
+</tr><tr>
+<td>
+YETI_LOGISTIC_LOSS<a id="YETI_LOGISTIC_LOSS"></a>
+</td>
+<td>
+`'yeti_logistic_loss'`
+</td>
 </tr>
 </table>
diff --git a/docs/api_docs/python/tfr/keras/losses/SigmoidCrossEntropyLoss.md b/docs/api_docs/python/tfr/keras/losses/SigmoidCrossEntropyLoss.md
index 768ad2d..ee1cb37 100644
--- a/docs/api_docs/python/tfr/keras/losses/SigmoidCrossEntropyLoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/SigmoidCrossEntropyLoss.md
@@ -15,7 +15,7 @@ description: Computes the Sigmoid cross-entropy loss between y_true and y_pred.
 
 <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#L1225-L1272">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L1371-L1418">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -85,13 +85,14 @@ $$
 </td>
 <td>
 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 under a
-`tf.distribute.Strategy`, except via `Model.compile()` and
-`Model.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.
+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 under a `tf.distribute.Strategy`, except via
+`Model.compile()` and `Model.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>
@@ -145,7 +146,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L239-L242">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L266-L269">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -156,7 +157,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L221-L229">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L248-L256">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/SoftmaxLoss.md b/docs/api_docs/python/tfr/keras/losses/SoftmaxLoss.md
index 8f25a30..e05b23f 100644
--- a/docs/api_docs/python/tfr/keras/losses/SoftmaxLoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/SoftmaxLoss.md
@@ -15,7 +15,7 @@ description: Computes Softmax cross-entropy loss between y_true and y_pred.
 
 <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#L595-L669">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L741-L815">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -126,7 +126,7 @@ False, this loss will accept dense tensors.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L585-L592">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L728-L738">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -167,7 +167,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L575-L583">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L718-L726">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -178,7 +178,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L661-L669">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L807-L815">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/UniqueSoftmaxLoss.md b/docs/api_docs/python/tfr/keras/losses/UniqueSoftmaxLoss.md
index 9b8bfc8..7e1d7ec 100644
--- a/docs/api_docs/python/tfr/keras/losses/UniqueSoftmaxLoss.md
+++ b/docs/api_docs/python/tfr/keras/losses/UniqueSoftmaxLoss.md
@@ -16,7 +16,7 @@ y_pred.
 
 <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#L672-L731">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L818-L877">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -104,13 +104,14 @@ $$
 </td>
 <td>
 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 under a
-`tf.distribute.Strategy`, except via `Model.compile()` and
-`Model.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.
+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 under a `tf.distribute.Strategy`, except via
+`Model.compile()` and `Model.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>
@@ -126,7 +127,7 @@ Optional name for the instance.
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L585-L592">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L728-L738">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -167,7 +168,7 @@ A `Loss` instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L575-L583">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L718-L726">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -178,7 +179,7 @@ Returns the config dictionary for a `Loss` instance.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L221-L229">View
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L248-L256">View
 source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/losses/YetiDCGLambdaWeight.md b/docs/api_docs/python/tfr/keras/losses/YetiDCGLambdaWeight.md
new file mode 100644
index 0000000..eb0cba8
--- /dev/null
+++ b/docs/api_docs/python/tfr/keras/losses/YetiDCGLambdaWeight.md
@@ -0,0 +1,110 @@
+description: Keras serializable class for YetiDCGLambdaWeight.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.losses.YetiDCGLambdaWeight" />
+<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.YetiDCGLambdaWeight
+
+<!-- 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#L156-L178">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Keras serializable class for YetiDCGLambdaWeight.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.losses.YetiDCGLambdaWeight(
+    topn: Optional[int] = None,
+    gain_fn: Optional[<a href="../../../tfr/keras/utils/GainFunction.md"><code>tfr.keras.utils.GainFunction</code></a>] = None,
+    rank_discount_fn: Optional[<a href="../../../tfr/keras/utils/GainFunction.md"><code>tfr.keras.utils.GainFunction</code></a>] = None,
+    normalized: bool = 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">Args</h2></th></tr>
+
+<tr>
+<td>
+`topn`<a id="topn"></a>
+</td>
+<td>
+(int) The topn for the DCG metric.
+</td>
+</tr><tr>
+<td>
+`gain_fn`<a id="gain_fn"></a>
+</td>
+<td>
+(function) Transforms labels.
+</td>
+</tr><tr>
+<td>
+`rank_discount_fn`<a id="rank_discount_fn"></a>
+</td>
+<td>
+(function) The rank discount function.
+</td>
+</tr><tr>
+<td>
+`normalized`<a id="normalized"></a>
+</td>
+<td>
+(bool) If True, normalize weight by the max DCG.
+</td>
+</tr>
+</table>
+
+## Methods
+
+<h3 id="get_config"><code>get_config</code></h3>
+
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L172-L178">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_config() -> Dict[str, Any]
+</code></pre>
+
+<h3 id="individual_weights"><code>individual_weights</code></h3>
+
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L267-L282">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" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/losses_impl.py#L386-L393">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>pair_weights(
+    labels: tf.Tensor, ranks: tf.Tensor
+) -> tf.Tensor
+</code></pre>
+
+See `_LambdaWeight`.
diff --git a/docs/api_docs/python/tfr/keras/losses/YetiLogisticLoss.md b/docs/api_docs/python/tfr/keras/losses/YetiLogisticLoss.md
new file mode 100644
index 0000000..ab094cd
--- /dev/null
+++ b/docs/api_docs/python/tfr/keras/losses/YetiLogisticLoss.md
@@ -0,0 +1,204 @@
+description: Computes Yeti logistic loss between y_true and y_pred.
+
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.losses.YetiLogisticLoss" />
+<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.YetiLogisticLoss
+
+<!-- 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#L591-L701">
+    <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
+    View source on GitHub
+  </a>
+</td>
+</table>
+
+Computes Yeti logistic loss between `y_true` and `y_pred`.
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>tfr.keras.losses.YetiLogisticLoss(
+    reduction: tf.losses.Reduction = tf.losses.Reduction.AUTO,
+    name: Optional[str] = None,
+    lambda_weight: Optional[<a href="../../../tfr/keras/losses/YetiDCGLambdaWeight.md"><code>tfr.keras.losses.YetiDCGLambdaWeight</code></a>] = None,
+    temperature: float = 0.1,
+    sample_size: int = 8,
+    gumbel_temperature: float = 1.0,
+    seed: Optional[int] = None,
+    ragged: bool = False
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
+
+Adapted to neural network models from the Yeti loss implemenation for GBDT in
+([Lyzhin et al, 2022][lyzhin2022]).
+
+In this code base, we support Yeti loss with the DCG lambda weight option. The
+default uses the YetiDCGLambdaWeight with default settings. To customize, please
+set the lambda_weight to YetiDCGLambdaWeight.
+
+For each list of scores `s` in `y_pred` and list of labels `y` in `y_true`:
+
+```
+loss = sum_a sum_i I[y_i > y_{i\pm 1}] * log(1 + exp(-(s^a_i - s^a_{i\pm 1})))
+```
+
+where `s^a_i = s_i + gumbel(0, 1)^a`
+
+#### Standalone usage:
+
+```
+>>> y_true = [[1., 0.]]
+>>> y_pred = [[0.6, 0.8]]
+>>> loss = tfr.keras.losses.YetiLogisticLoss(sample_size=2, seed=1)
+>>> loss(y_true, y_pred).numpy()
+0.90761846
+```
+
+```
+>>> # Using ragged tensors
+>>> y_true = tf.ragged.constant([[1., 0.], [0., 1., 0.]])
+>>> y_pred = tf.ragged.constant([[0.6, 0.8], [0.5, 0.8, 0.4]])
+>>> loss = tfr.keras.losses.YetiLogisticLoss(seed=1, ragged=True)
+>>> loss(y_true, y_pred).numpy()
+0.43420443
+```
+
+Usage with the `compile()` API:
+
+```python
+model.compile(optimizer='sgd', loss=tfr.keras.losses.YetiLogisticLoss())
+```
+
+#### Definition:
+
+$$
+\mathcal{L}(\{y\}, \{s\}) =
+\sum_a \sum_i \sum_{j=i\pm 1}I[y_i > y_j] \log(1 + \exp(-(s^a_i - s^a_j)))
+$$
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2"><h2 class="add-link">References</h2></th></tr>
+<tr class="alt">
+<td colspan="2">
+- [Which Tricks are Important for Learning to Rank?, Lyzhin et al, 2022][lyzhin2022]
+</td>
+</tr>
+
+</table>
+
+[lyzhin2022]: https://arxiv.org/abs/2204.01500
+
+<!-- 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`<a id="reduction"></a>
+</td>
+<td>
+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 under a `tf.distribute.Strategy`, except via
+`Model.compile()` and `Model.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`<a id="name"></a>
+</td>
+<td>
+Optional name for the instance.
+</td>
+</tr>
+</table>
+
+## Methods
+
+<h3 id="from_config"><code>from_config</code></h3>
+
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L296-L306">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>@classmethod</code>
+<code>from_config(
+    config, custom_objects=None
+)
+</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" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L680-L687">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>get_config() -> Dict[str, Any]
+</code></pre>
+
+Returns the config dictionary for a `Loss` instance.
+
+<h3 id="__call__"><code>__call__</code></h3>
+
+<a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L689-L701">View
+source</a>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>__call__(
+    y_true: <a href="../../../tfr/keras/model/TensorLike.md"><code>tfr.keras.model.TensorLike</code></a>,
+    y_pred: <a href="../../../tfr/keras/model/TensorLike.md"><code>tfr.keras.model.TensorLike</code></a>,
+    sample_weight: Optional[utils.TensorLike] = None
+) -> tf.Tensor
+</code></pre>
+
+See _RankingLoss.
diff --git a/docs/api_docs/python/tfr/keras/losses/get.md b/docs/api_docs/python/tfr/keras/losses/get.md
index f1ac68a..29c5d0e 100644
--- a/docs/api_docs/python/tfr/keras/losses/get.md
+++ b/docs/api_docs/python/tfr/keras/losses/get.md
@@ -11,7 +11,7 @@ description: Factory method to get a ranking loss class.
 
 <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#L35-L93">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/losses.py#L36-L95">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/docs/api_docs/python/tfr/keras/metrics/ARPMetric.md b/docs/api_docs/python/tfr/keras/metrics/ARPMetric.md
index 70738bb..c9d9334 100644
--- a/docs/api_docs/python/tfr/keras/metrics/ARPMetric.md
+++ b/docs/api_docs/python/tfr/keras/metrics/ARPMetric.md
@@ -7,7 +7,6 @@ description: Average relevance position (ARP).
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,9 +16,11 @@ description: Average relevance position (ARP).
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
 <meta itemprop="property" content="merge_state"/>
 <meta itemprop="property" content="reset_state"/>
 <meta itemprop="property" content="result"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="update_state"/>
 <meta itemprop="property" content="with_name_scope"/>
@@ -192,20 +193,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -352,89 +341,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -477,6 +383,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -653,6 +582,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L181-L186">View
@@ -714,6 +667,35 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
 <h3 id="merge_state"><code>merge_state</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -814,6 +796,35 @@ A scalar tensor, or a dictionary of scalar tensors.
 
 </table>
 
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/metrics/AlphaDCGMetric.md b/docs/api_docs/python/tfr/keras/metrics/AlphaDCGMetric.md
index 6cc6598..7ed1ff5 100644
--- a/docs/api_docs/python/tfr/keras/metrics/AlphaDCGMetric.md
+++ b/docs/api_docs/python/tfr/keras/metrics/AlphaDCGMetric.md
@@ -7,7 +7,6 @@ description: Alpha discounted cumulative gain (alphaDCG).
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,9 +16,11 @@ description: Alpha discounted cumulative gain (alphaDCG).
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
 <meta itemprop="property" content="merge_state"/>
 <meta itemprop="property" content="reset_state"/>
 <meta itemprop="property" content="result"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="update_state"/>
 <meta itemprop="property" content="with_name_scope"/>
@@ -309,20 +310,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -469,89 +458,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -594,6 +500,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -770,6 +699,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L987-L995">View
@@ -831,6 +784,35 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
 <h3 id="merge_state"><code>merge_state</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -931,6 +913,35 @@ A scalar tensor, or a dictionary of scalar tensors.
 
 </table>
 
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/metrics/DCGMetric.md b/docs/api_docs/python/tfr/keras/metrics/DCGMetric.md
index 84ae3da..5b18b83 100644
--- a/docs/api_docs/python/tfr/keras/metrics/DCGMetric.md
+++ b/docs/api_docs/python/tfr/keras/metrics/DCGMetric.md
@@ -7,7 +7,6 @@ description: Discounted cumulative gain (DCG).
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,9 +16,11 @@ description: Discounted cumulative gain (DCG).
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
 <meta itemprop="property" content="merge_state"/>
 <meta itemprop="property" content="reset_state"/>
 <meta itemprop="property" content="result"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="update_state"/>
 <meta itemprop="property" content="with_name_scope"/>
@@ -222,20 +223,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -382,89 +371,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -507,6 +413,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -683,6 +612,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L860-L868">View
@@ -744,6 +697,35 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
 <h3 id="merge_state"><code>merge_state</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -844,6 +826,35 @@ A scalar tensor, or a dictionary of scalar tensors.
 
 </table>
 
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/metrics/HitsMetric.md b/docs/api_docs/python/tfr/keras/metrics/HitsMetric.md
index 23d59c1..643c92f 100644
--- a/docs/api_docs/python/tfr/keras/metrics/HitsMetric.md
+++ b/docs/api_docs/python/tfr/keras/metrics/HitsMetric.md
@@ -7,7 +7,6 @@ description: Hits@k metric.
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,9 +16,11 @@ description: Hits@k metric.
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
 <meta itemprop="property" content="merge_state"/>
 <meta itemprop="property" content="reset_state"/>
 <meta itemprop="property" content="result"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="update_state"/>
 <meta itemprop="property" content="with_name_scope"/>
@@ -202,20 +203,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -362,89 +351,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -487,6 +393,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -663,6 +592,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L312-L317">View
@@ -724,6 +677,35 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
 <h3 id="merge_state"><code>merge_state</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -824,6 +806,35 @@ A scalar tensor, or a dictionary of scalar tensors.
 
 </table>
 
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/metrics/MRRMetric.md b/docs/api_docs/python/tfr/keras/metrics/MRRMetric.md
index 4bd26c1..41b4926 100644
--- a/docs/api_docs/python/tfr/keras/metrics/MRRMetric.md
+++ b/docs/api_docs/python/tfr/keras/metrics/MRRMetric.md
@@ -7,7 +7,6 @@ description: Mean reciprocal rank (MRR).
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,9 +16,11 @@ description: Mean reciprocal rank (MRR).
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
 <meta itemprop="property" content="merge_state"/>
 <meta itemprop="property" content="reset_state"/>
 <meta itemprop="property" content="result"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="update_state"/>
 <meta itemprop="property" content="with_name_scope"/>
@@ -201,20 +202,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -361,89 +350,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -486,6 +392,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -662,6 +591,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L246-L251">View
@@ -723,6 +676,35 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
 <h3 id="merge_state"><code>merge_state</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -823,6 +805,35 @@ A scalar tensor, or a dictionary of scalar tensors.
 
 </table>
 
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/metrics/MeanAveragePrecisionMetric.md b/docs/api_docs/python/tfr/keras/metrics/MeanAveragePrecisionMetric.md
index 7b3bd63..b1e4671 100644
--- a/docs/api_docs/python/tfr/keras/metrics/MeanAveragePrecisionMetric.md
+++ b/docs/api_docs/python/tfr/keras/metrics/MeanAveragePrecisionMetric.md
@@ -7,7 +7,6 @@ description: Mean average precision (MAP).
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,9 +16,11 @@ description: Mean average precision (MAP).
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
 <meta itemprop="property" content="merge_state"/>
 <meta itemprop="property" content="reset_state"/>
 <meta itemprop="property" content="result"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="update_state"/>
 <meta itemprop="property" content="with_name_scope"/>
@@ -209,20 +210,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -369,89 +358,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -494,6 +400,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -670,6 +599,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L686-L692">View
@@ -731,6 +684,35 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
 <h3 id="merge_state"><code>merge_state</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -831,6 +813,35 @@ A scalar tensor, or a dictionary of scalar tensors.
 
 </table>
 
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/metrics/NDCGMetric.md b/docs/api_docs/python/tfr/keras/metrics/NDCGMetric.md
index 92f088a..c9c52ad 100644
--- a/docs/api_docs/python/tfr/keras/metrics/NDCGMetric.md
+++ b/docs/api_docs/python/tfr/keras/metrics/NDCGMetric.md
@@ -7,7 +7,6 @@ description: Normalized discounted cumulative gain (NDCG).
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,9 +16,11 @@ description: Normalized discounted cumulative gain (NDCG).
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
 <meta itemprop="property" content="merge_state"/>
 <meta itemprop="property" content="reset_state"/>
 <meta itemprop="property" content="result"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="update_state"/>
 <meta itemprop="property" content="with_name_scope"/>
@@ -227,20 +228,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -387,89 +376,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -512,6 +418,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -688,6 +617,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L774-L782">View
@@ -749,6 +702,35 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
 <h3 id="merge_state"><code>merge_state</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -849,6 +831,35 @@ A scalar tensor, or a dictionary of scalar tensors.
 
 </table>
 
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/metrics/OPAMetric.md b/docs/api_docs/python/tfr/keras/metrics/OPAMetric.md
index 9977e47..7719964 100644
--- a/docs/api_docs/python/tfr/keras/metrics/OPAMetric.md
+++ b/docs/api_docs/python/tfr/keras/metrics/OPAMetric.md
@@ -7,7 +7,6 @@ description: Ordered pair accuracy (OPA).
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,9 +16,11 @@ description: Ordered pair accuracy (OPA).
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
 <meta itemprop="property" content="merge_state"/>
 <meta itemprop="property" content="reset_state"/>
 <meta itemprop="property" content="result"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="update_state"/>
 <meta itemprop="property" content="with_name_scope"/>
@@ -200,20 +201,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -360,89 +349,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -485,6 +391,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -661,6 +590,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L181-L186">View
@@ -722,6 +675,35 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
 <h3 id="merge_state"><code>merge_state</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -822,6 +804,35 @@ A scalar tensor, or a dictionary of scalar tensors.
 
 </table>
 
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/metrics/PrecisionIAMetric.md b/docs/api_docs/python/tfr/keras/metrics/PrecisionIAMetric.md
index 8ab15f2..00f6204 100644
--- a/docs/api_docs/python/tfr/keras/metrics/PrecisionIAMetric.md
+++ b/docs/api_docs/python/tfr/keras/metrics/PrecisionIAMetric.md
@@ -7,7 +7,6 @@ description: Precision-IA@k (Pre-IA@k).
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,9 +16,11 @@ description: Precision-IA@k (Pre-IA@k).
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
 <meta itemprop="property" content="merge_state"/>
 <meta itemprop="property" content="reset_state"/>
 <meta itemprop="property" content="result"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="update_state"/>
 <meta itemprop="property" content="with_name_scope"/>
@@ -266,20 +267,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -426,89 +415,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -551,6 +457,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -727,6 +656,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L606-L611">View
@@ -788,6 +741,35 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
 <h3 id="merge_state"><code>merge_state</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -888,6 +870,35 @@ A scalar tensor, or a dictionary of scalar tensors.
 
 </table>
 
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/metrics/PrecisionMetric.md b/docs/api_docs/python/tfr/keras/metrics/PrecisionMetric.md
index acc0da2..83f97d1 100644
--- a/docs/api_docs/python/tfr/keras/metrics/PrecisionMetric.md
+++ b/docs/api_docs/python/tfr/keras/metrics/PrecisionMetric.md
@@ -7,7 +7,6 @@ description: Precision@k (P@k).
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,9 +16,11 @@ description: Precision@k (P@k).
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
 <meta itemprop="property" content="merge_state"/>
 <meta itemprop="property" content="reset_state"/>
 <meta itemprop="property" content="result"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="update_state"/>
 <meta itemprop="property" content="with_name_scope"/>
@@ -204,20 +205,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -364,89 +353,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -489,6 +395,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -665,6 +594,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L435-L440">View
@@ -726,6 +679,35 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
 <h3 id="merge_state"><code>merge_state</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -826,6 +808,35 @@ A scalar tensor, or a dictionary of scalar tensors.
 
 </table>
 
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/metrics/RecallMetric.md b/docs/api_docs/python/tfr/keras/metrics/RecallMetric.md
index 088eae0..47f76a4 100644
--- a/docs/api_docs/python/tfr/keras/metrics/RecallMetric.md
+++ b/docs/api_docs/python/tfr/keras/metrics/RecallMetric.md
@@ -7,7 +7,6 @@ description: Recall@k (R@k).
 <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="build_from_config"/>
 <meta itemprop="property" content="compute_mask"/>
@@ -17,9 +16,11 @@ description: Recall@k (R@k).
 <meta itemprop="property" content="get_build_config"/>
 <meta itemprop="property" content="get_config"/>
 <meta itemprop="property" content="get_weights"/>
+<meta itemprop="property" content="load_own_variables"/>
 <meta itemprop="property" content="merge_state"/>
 <meta itemprop="property" content="reset_state"/>
 <meta itemprop="property" content="result"/>
+<meta itemprop="property" content="save_own_variables"/>
 <meta itemprop="property" content="set_weights"/>
 <meta itemprop="property" content="update_state"/>
 <meta itemprop="property" content="with_name_scope"/>
@@ -204,20 +205,8 @@ gradients back to the corresponding variables.
 ```
 
 </td> </tr><tr> <td> `metrics`<a id="metrics"></a> </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`<a id="name"></a> </td> <td> Name of the layer
-(string), set in the constructor. </td> </tr><tr> <td>
+attached to the layer. </td> </tr><tr> <td> `name`<a id="name"></a> </td> <td>
+Name of the layer (string), set in the constructor. </td> </tr><tr> <td>
 `name_scope`<a id="name_scope"></a> </td> <td> Returns a `tf.name_scope`
 instance for this class. </td> </tr><tr> <td>
 `non_trainable_weights`<a id="non_trainable_weights"></a> </td> <td> List of all
@@ -364,89 +353,6 @@ Used for backwards compatibility only.
 </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(inputs))
-    self.add_metric(tf.reduce_sum(inputs), 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">
@@ -489,6 +395,29 @@ Instance of `TensorShape`, or list of instances of
 )
 </code></pre>
 
+Builds the layer's states with the supplied config dict.
+
+By default, this method calls the `build(config["input_shape"])` method, which
+creates weights based on the layer's input shape in the supplied config. If your
+config contains other information needed to load the layer's state, you should
+override this method.
+
+<!-- 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>
+Dict containing the input shape associated with this layer.
+</td>
+</tr>
+</table>
+
 <h3 id="compute_mask"><code>compute_mask</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -665,6 +594,30 @@ A layer instance.
 <code>get_build_config()
 </code></pre>
 
+Returns a dictionary with the layer's input shape.
+
+This method returns a config dict that can be used by
+`build_from_config(config)` to create all states (e.g. Variables and Lookup
+tables) needed by the layer.
+
+By default, the config only contains the input shape that the layer was built
+with. If you're writing a custom layer that creates state in an unusual way, you
+should override this method to make sure this state is already created when
+Keras attempts to load its value upon model loading.
+
+<!-- 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 containing the input shape associated with the layer.
+</td>
+</tr>
+
+</table>
+
 <h3 id="get_config"><code>get_config</code></h3>
 
 <a target="_blank" class="external" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/metrics.py#L511-L516">View
@@ -726,6 +679,35 @@ Weights values as a list of NumPy arrays.
 
 </table>
 
+<h3 id="load_own_variables"><code>load_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>load_own_variables(
+    store
+)
+</code></pre>
+
+Loads the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is loaded upon calling `keras.models.load_model()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict from which the state of the model will be loaded.
+</td>
+</tr>
+</table>
+
 <h3 id="merge_state"><code>merge_state</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
@@ -826,6 +808,35 @@ A scalar tensor, or a dictionary of scalar tensors.
 
 </table>
 
+<h3 id="save_own_variables"><code>save_own_variables</code></h3>
+
+<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
+<code>save_own_variables(
+    store
+)
+</code></pre>
+
+Saves the state of the layer.
+
+You can override this method to take full control of how the state of the layer
+is saved upon calling `model.save()`.
+
+<!-- Tabular view -->
+
+ <table class="responsive fixed orange">
+<colgroup><col width="214px"><col></colgroup>
+<tr><th colspan="2">Args</th></tr>
+
+<tr>
+<td>
+`store`
+</td>
+<td>
+Dict where the state of the model will be saved.
+</td>
+</tr>
+</table>
+
 <h3 id="set_weights"><code>set_weights</code></h3>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
diff --git a/docs/api_docs/python/tfr/keras/pipeline/DatasetHparams/dataset_reader.md b/docs/api_docs/python/tfr/keras/pipeline/DatasetHparams/dataset_reader.md
index 87f5fc7..052cfe8 100644
--- a/docs/api_docs/python/tfr/keras/pipeline/DatasetHparams/dataset_reader.md
+++ b/docs/api_docs/python/tfr/keras/pipeline/DatasetHparams/dataset_reader.md
@@ -2731,7 +2731,6 @@ Optional. A name for the tf.data transformation.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -2788,7 +2787,6 @@ TensorShape([None])
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Args</th></tr>
@@ -2831,7 +2829,6 @@ their row_splits dtype changed.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -3567,6 +3564,20 @@ users should not set the `checkpoint` argument in `checkpoint_args`.
 </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 operation which when executed performs the save. When writing
+checkpoints, returns None. The return value is useful in unit tests.
+</td>
+</tr>
+
+</table>
+
 <!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
@@ -3821,6 +3832,24 @@ list(dataset.as_numpy_iterator())
 # [1, 0, 2]
 ```
 
+### Fully shuffling all the data
+
+To shuffle an entire dataset, set `buffer_size=dataset.cardinality(). This is
+equivalent to setting the`buffer_size` equal to the number of elements in the
+dataset, resulting in uniform shuffle.
+
+Note: `shuffle(dataset.cardinality())` loads the full dataset into memory so
+that it can be shuffled. This will cause a memory overflow (OOM) error if the
+dataset is too large, so full-shuffle should only be used for datasets that are
+known to fit in the memory, such as datasets of filenames or other small
+datasets.
+
+```python
+dataset = tf.data.Dataset.range(20)
+dataset = dataset.shuffle(dataset.cardinality())
+# [18, 4, 9, 2, 17, 8, 5, 10, 0, 6, 16, 3, 19, 7, 14, 11, 15, 13, 12, 1]
+```
+
 <!-- Tabular view -->
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
@@ -3832,7 +3861,9 @@ list(dataset.as_numpy_iterator())
 </td>
 <td>
 A `tf.int64` scalar `tf.Tensor`, representing the number of
-elements from this dataset from which the new dataset will sample.
+elements from this dataset from which the new dataset will sample. To
+uniformly shuffle the entire dataset, use
+`buffer_size=dataset.cardinality()`.
 </td>
 </tr><tr>
 <td>
@@ -4048,7 +4079,6 @@ a snapshot.
 </table>
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Returns</th></tr>
@@ -4096,7 +4126,6 @@ a.apply(tf.data.experimental.dense_to_sparse_batch(
 ```
 
 <!-- Tabular view -->
-
  <table class="responsive fixed orange">
 <colgroup><col width="214px"><col></colgroup>
 <tr><th colspan="2">Args</th></tr>
@@ -4645,7 +4674,7 @@ when an option is set more than once to a non-default value
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@staticmethod</code>
 <code>zip(
-    datasets, name=None
+    *args, datasets=None, name=None
 )
 </code></pre>
 
@@ -4657,14 +4686,14 @@ structure of `Dataset` objects. The supported nesting mechanisms are documented
 [here](https://www.tensorflow.org/guide/data#dataset_structure).
 
 ```
->>> # The nested structure of the `datasets` argument determines the
->>> # structure of elements in the resulting dataset.
+>>> # The datasets or nested structure of datasets `*args` argument
+>>> # determines the structure of elements in the resulting dataset.
 >>> a = tf.data.Dataset.range(1, 4)  # ==> [ 1, 2, 3 ]
 >>> b = tf.data.Dataset.range(4, 7)  # ==> [ 4, 5, 6 ]
->>> ds = tf.data.Dataset.zip((a, b))
+>>> ds = tf.data.Dataset.zip(a, b)
 >>> list(ds.as_numpy_iterator())
 [(1, 4), (2, 5), (3, 6)]
->>> ds = tf.data.Dataset.zip((b, a))
+>>> ds = tf.data.Dataset.zip(b, a)
 >>> list(ds.as_numpy_iterator())
 [(4, 1), (5, 2), (6, 3)]
 >>>
@@ -4672,7 +4701,7 @@ structure of `Dataset` objects. The supported nesting mechanisms are documented
 >>> c = tf.data.Dataset.range(7, 13).batch(2)  # ==> [ [7, 8],
 ...                                            #       [9, 10],
 ...                                            #       [11, 12] ]
->>> ds = tf.data.Dataset.zip((a, b, c))
+>>> ds = tf.data.Dataset.zip(a, b, c)
 >>> for element in ds.as_numpy_iterator():
 ...   print(element)
 (1, 4, array([7, 8]))
@@ -4682,7 +4711,7 @@ structure of `Dataset` objects. The supported nesting mechanisms are documented
 >>> # The number of elements in the resulting dataset is the same as
 >>> # the size of the smallest dataset in `datasets`.
 >>> d = tf.data.Dataset.range(13, 15)  # ==> [ 13, 14 ]
->>> ds = tf.data.Dataset.zip((a, d))
+>>> ds = tf.data.Dataset.zip(a, d)
 >>> list(ds.as_numpy_iterator())
 [(1, 13), (2, 14)]
 ```
@@ -4694,10 +4723,20 @@ structure of `Dataset` objects. The supported nesting mechanisms are documented
 
 <tr>
 <td>
+`*args`
+</td>
+<td>
+Datasets or nested structures of datasets to zip together. This
+can't be set if `datasets` is set.
+</td>
+</tr><tr>
+<td>
 `datasets`
 </td>
 <td>
-A (nested) structure of datasets.
+A (nested) structure of datasets. This can't be set if `*args`
+is set. Note that this exists only for backwards compatibility and it is
+preferred to use *args.
 </td>
 </tr><tr>
 <td>
diff --git a/docs/api_docs/python/tfr/keras/utils.md b/docs/api_docs/python/tfr/keras/utils.md
index fc0c4f2..5e0f32a 100644
--- a/docs/api_docs/python/tfr/keras/utils.md
+++ b/docs/api_docs/python/tfr/keras/utils.md
@@ -22,6 +22,8 @@ Utils for tfr.keras.
 
 ## Functions
 
+[`deserialize_keras_object(...)`](../../tfr/keras/utils/deserialize_keras_object.md)
+
 [`identity(...)`](../../tfr/keras/utils/identity.md): Identity function that
 returns the input label.
 
@@ -37,6 +39,8 @@ Computes whether label is greater or equal to 1.
 [`pow_minus_1(...)`](../../tfr/keras/utils/pow_minus_1.md): Computes `2**x - 1`
 element-wise for each label.
 
+[`serialize_keras_object(...)`](../../tfr/keras/utils/serialize_keras_object.md)
+
 [`symmetric_log1p(...)`](../../tfr/keras/utils/symmetric_log1p.md): Computes
 `sign(x) * log(1 + sign(x))`.
 
diff --git a/docs/api_docs/python/tfr/keras/utils/GainFunction.md b/docs/api_docs/python/tfr/keras/utils/GainFunction.md
index 601a3fb..6a57b20 100644
--- a/docs/api_docs/python/tfr/keras/utils/GainFunction.md
+++ b/docs/api_docs/python/tfr/keras/utils/GainFunction.md
@@ -6,7 +6,6 @@
 # tfr.keras.utils.GainFunction
 
 <!-- Insert buttons and diff -->
-
 This symbol is a **type alias**.
 
 #### Source:
diff --git a/docs/api_docs/python/tfr/keras/utils/deserialize_keras_object.md b/docs/api_docs/python/tfr/keras/utils/deserialize_keras_object.md
new file mode 100644
index 0000000..410d34c
--- /dev/null
+++ b/docs/api_docs/python/tfr/keras/utils/deserialize_keras_object.md
@@ -0,0 +1,28 @@
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.utils.deserialize_keras_object" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# tfr.keras.utils.deserialize_keras_object
+
+<!-- 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-L30">
+    <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.deserialize_keras_object(
+    config,
+    module_objects=None,
+    custom_objects=None,
+    printable_module_name=None
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
diff --git a/docs/api_docs/python/tfr/keras/utils/identity.md b/docs/api_docs/python/tfr/keras/utils/identity.md
index 38b84f4..c5f7342 100644
--- a/docs/api_docs/python/tfr/keras/utils/identity.md
+++ b/docs/api_docs/python/tfr/keras/utils/identity.md
@@ -11,7 +11,7 @@ description: Identity function that returns the input label.
 
 <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#L16-L27">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/utils.py#L36-L47">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/docs/api_docs/python/tfr/keras/utils/inverse.md b/docs/api_docs/python/tfr/keras/utils/inverse.md
index df6a591..c489f6d 100644
--- a/docs/api_docs/python/tfr/keras/utils/inverse.md
+++ b/docs/api_docs/python/tfr/keras/utils/inverse.md
@@ -11,7 +11,7 @@ description: Computes the inverse of input rank.
 
 <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-L41">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/utils.py#L50-L61">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/docs/api_docs/python/tfr/keras/utils/is_greater_equal_1.md b/docs/api_docs/python/tfr/keras/utils/is_greater_equal_1.md
index aaca9fa..03e8aab 100644
--- a/docs/api_docs/python/tfr/keras/utils/is_greater_equal_1.md
+++ b/docs/api_docs/python/tfr/keras/utils/is_greater_equal_1.md
@@ -11,7 +11,7 @@ description: Computes whether label is greater or equal to 1.
 
 <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#L76-L87">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/utils.py#L96-L107">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/docs/api_docs/python/tfr/keras/utils/log2_inverse.md b/docs/api_docs/python/tfr/keras/utils/log2_inverse.md
index 2d83921..1d13b21 100644
--- a/docs/api_docs/python/tfr/keras/utils/log2_inverse.md
+++ b/docs/api_docs/python/tfr/keras/utils/log2_inverse.md
@@ -11,7 +11,7 @@ description: Computes 1./log2(1+x) element-wise for each label.
 
 <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#L60-L73">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/utils.py#L80-L93">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/docs/api_docs/python/tfr/keras/utils/pow_minus_1.md b/docs/api_docs/python/tfr/keras/utils/pow_minus_1.md
index ff94fc9..7e37be8 100644
--- a/docs/api_docs/python/tfr/keras/utils/pow_minus_1.md
+++ b/docs/api_docs/python/tfr/keras/utils/pow_minus_1.md
@@ -11,7 +11,7 @@ description: Computes 2**x - 1 element-wise for each label.
 
 <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#L44-L57">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/utils.py#L64-L77">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/docs/api_docs/python/tfr/keras/utils/serialize_keras_object.md b/docs/api_docs/python/tfr/keras/utils/serialize_keras_object.md
new file mode 100644
index 0000000..f7f670f
--- /dev/null
+++ b/docs/api_docs/python/tfr/keras/utils/serialize_keras_object.md
@@ -0,0 +1,25 @@
+<div itemscope itemtype="http://developers.google.com/ReferenceObject">
+<meta itemprop="name" content="tfr.keras.utils.serialize_keras_object" />
+<meta itemprop="path" content="Stable" />
+</div>
+
+# tfr.keras.utils.serialize_keras_object
+
+<!-- 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#L13-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.serialize_keras_object(
+    obj
+)
+</code></pre>
+
+<!-- Placeholder for "Used in" -->
diff --git a/docs/api_docs/python/tfr/keras/utils/symmetric_log1p.md b/docs/api_docs/python/tfr/keras/utils/symmetric_log1p.md
index 8b584e5..7f4dede 100644
--- a/docs/api_docs/python/tfr/keras/utils/symmetric_log1p.md
+++ b/docs/api_docs/python/tfr/keras/utils/symmetric_log1p.md
@@ -11,7 +11,7 @@ description: Computes sign(x) * log(1 + sign(x)).
 
 <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#L90-L101">
+  <a target="_blank" href="https://github.com/tensorflow/ranking/tree/master/tensorflow_ranking/python/keras/utils.py#L110-L121">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
diff --git a/docs/api_docs/python/tfr/utils/LossFunction.md b/docs/api_docs/python/tfr/utils/LossFunction.md
index 257807a..701fa49 100644
--- a/docs/api_docs/python/tfr/utils/LossFunction.md
+++ b/docs/api_docs/python/tfr/utils/LossFunction.md
@@ -6,7 +6,6 @@
 # tfr.utils.LossFunction
 
 <!-- Insert buttons and diff -->
-
 This symbol is a **type alias**.
 
 #### Source:
diff --git a/docs/api_docs/python/tfr/utils/parse_keys_and_weights.md b/docs/api_docs/python/tfr/utils/parse_keys_and_weights.md
index e13b4b2..35c067a 100644
--- a/docs/api_docs/python/tfr/utils/parse_keys_and_weights.md
+++ b/docs/api_docs/python/tfr/utils/parse_keys_and_weights.md
@@ -32,7 +32,6 @@ This parse function will remove all spaces. Different keys are split by "," and
 then weight associated with key is split by ":".
 
 <!-- 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>
@@ -50,7 +49,6 @@ and weighted by the weights split by ":". For example, key =
 </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>