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

Reviewed By: jackh, #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved

Maniphest Tasks: T43586

Differential Revision: https://phabricator.sourcevertex.net/D49238

Author: georgepaw

tensorflow/python/ipu/keras/optimizers/ipu_wrappers.py

@@ -16,6 +16,7 @@
 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 @@ def compute_gradients(self,
 
     v = var_list if not self._model else self._model.trainable_weights
 
-    grads = self._optimizer.get_gradients(loss, v)
-    grads_and_vars = zip(grads, v)
+    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):