message: Make sure gradient accumulation calls the right functions
  Summary:
  Fix T43586
  TF2.4 Only
  Test Plan: CI
  Reviewers: jackh, samuelh, alfiee, #tensorflow, simonl, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved
  Maniphest Tasks: T43586
diff: Convenience wrappers for v2 optimizers
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  """
  +from tensorflow.python.keras.optimizer_v1 import TFOptimizer
  from tensorflow.python.keras.optimizer_v2.optimizer_v2 import OptimizerV2
  from tensorflow.python.training.optimizer import Optimizer
  @@ -169,8 +170,12 @@ class _KerasOptimizerWrapper(Optimizer):
  v = var_list if not self._model else self._model.trainable_weights
  + if isinstance(self._optimizer, TFOptimizer):
  + grads_and_vars = self._optimizer.get_grads(loss, v)
  + else:
  grads = self._optimizer.get_gradients(loss, v)
  grads_and_vars = zip(grads, v)
  +
  return list(map(self.preprocess_gradients, grads_and_vars))
  def apply_gradients(self, grads_and_vars, global_step=None, name=None):
  