Merged commit includes the following changes: (#7013)

253113801  by A. Unique TensorFlower<gardener@tensorflow.org>:

    Internal change

252697519  by dmchen<dmchen@google.com>:

    BERT SQuAD accuracy test

--
252663512  by A. Unique TensorFlower<gardener@tensorflow.org>:

    Internal change

--
252647871  by A. Unique TensorFlower<gardener@tensorflow.org>:

    Enable multi worker TPU training for BERT pretraining.

--

PiperOrigin-RevId: 253113801

official/bert/optimization.py

@@ -136,14 +136,14 @@ class AdamWeightDecay(tf.keras.optimizers.Adam):
 
   def _resource_apply_dense(self, grad, var):
     var_dtype = var.dtype.base_dtype
-    lr_t = self._decayed_lr(var_dtype)
+    lr_t = self._decayed_lr_t[var_dtype]
     with tf.control_dependencies([self._decay_weights_op(var, lr_t)]):
       return super(AdamWeightDecay, self)._resource_apply_dense(
           grad, var)
 
   def _resource_apply_sparse(self, grad, var, indices):
     var_dtype = var.dtype.base_dtype
-    lr_t = self._decayed_lr(var_dtype)
+    lr_t = self._decayed_lr_t[var_dtype]
     with tf.control_dependencies([self._decay_weights_op(var, lr_t)]):
       return super(AdamWeightDecay, self)._resource_apply_sparse(
           grad, var, indices)