Prevent None from becoming a control dependency

Summary:
Depending on the optimizer wrapped, None can be returned.

Fix T43585
TF2.4 Only

Test Plan: CI, added a test which failed before

Reviewers: jackh, alfiee, samuelh, #tensorflow, simonl, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, jakeh

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

Maniphest Tasks: T43585

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

Author: georgepaw

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

@@ -120,7 +120,8 @@ def _resource_apply_dense(self, grad, handle, apply_state):
     def resource_update_():
       updated_var = self._opt._resource_apply_dense(  # pylint: disable=protected-access
           acc_grad, acc_var, apply_state)
-      apply_grad_ops.append(updated_var)
+      if updated_var is not None:
+        apply_grad_ops.append(updated_var)
 
     return GradientAccumulationOptimizerV2.apply_gradient_accumulation(
         resource_update_,

tensorflow/python/ipu/ops/pipelining_ops.py

@@ -1009,15 +1009,15 @@ def resource_update_():
           apply_grads = opt.apply_gradients(accumulated_grads_and_vars,
                                             *apply_gradients_args,
                                             **apply_gradients_kwargs)
-          if apply_grads:
+          if apply_grads is not None:
             resource_update_ops.append(apply_grads)
 
         # Enqueue any accumulated outfeed data
         if outfeed_sinks:
           # Note: unpack if we're outfeeding loss.
           to_enqueue = outfeed_sinks[0] if outfeed_loss else outfeed_sinks
           enqueue = outfeed_queue.enqueue(to_enqueue)
-          if enqueue:
+          if enqueue is not None:
             resource_update_ops.append(enqueue)
 
       with ops.name_scope(name + "/WU") as scope:

tensorflow/python/ipu/optimizers/gradient_accumulation_optimizer.py

@@ -187,7 +187,8 @@ def apply_gradients(self, grads_and_vars, global_step=None, name=None):
     def resource_update_():
       apply_grads = self._opt.apply_gradients(accumulated_grads_and_vars,
                                               global_step, name)
-      apply_grad_ops.append(apply_grads)
+      if apply_grads is not None:
+        apply_grad_ops.append(apply_grads)
 
     return self.apply_gradient_accumulation(
         resource_update_, self._opt.get_name(), apply_grad_ops,

tensorflow/python/ipu/tests/keras/keras_gradient_accumulation_test.py

@@ -22,6 +22,7 @@
 from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.platform import test
+from tensorflow.python.training import gradient_descent
 
 
 def get_mnist_dataset(batch_size):
@@ -63,24 +64,28 @@ class KerasGradientAccumulationTest(test.TestCase, parameterized.TestCase):
   TESTCASES = [{
       "testcase_name": "sequential",
       "model_fn": simple_sequential_model,
-      "replication_factor": 1
+      "replication_factor": 1,
+      "optimizer": "sgd"
   }, {
       "testcase_name": "sequential_replicated",
       "model_fn": simple_sequential_model,
-      "replication_factor": 2
+      "replication_factor": 2,
+      "optimizer": "sgd"
   }, {
       "testcase_name": "functional",
       "model_fn": simple_functional_model,
-      "replication_factor": 1
+      "replication_factor": 1,
+      "optimizer": "sgd"
   }, {
       "testcase_name": "functional_replicated",
       "model_fn": simple_functional_model,
-      "replication_factor": 2
+      "replication_factor": 2,
+      "optimizer": gradient_descent.GradientDescentOptimizer(0.001)
   }]
 
   @parameterized.named_parameters(*TESTCASES)
   @test_util.run_v2_only
-  def testModels(self, model_fn, replication_factor):
+  def testModels(self, model_fn, replication_factor, optimizer):
     tu.skip_if_not_enough_ipus(self, replication_factor)
 
     cfg = ipu.config.IPUConfig()
@@ -98,8 +103,7 @@ def testModels(self, model_fn, replication_factor):
     # Run on CPU - simulate gradient accumulation by just using a bigger batch
     # size but less steps per epoch.
     m = model_fn()
-    m.compile(optimizer='sgd',
-              loss=keras.losses.SparseCategoricalCrossentropy())
+    m.compile(optimizer, loss=keras.losses.SparseCategoricalCrossentropy())
     m.fit(get_mnist_dataset(batch_size * gradient_accumulation_steps),
           steps_per_epoch=steps_per_epoch // gradient_accumulation_steps,
           epochs=epochs)
@@ -108,7 +112,7 @@ def testModels(self, model_fn, replication_factor):
     strategy = ipu.ipu_strategy.IPUStrategyV1()
     with strategy.scope():
       m = model_fn()
-      m.compile(optimizer='sgd',
+      m.compile(optimizer,
                 loss=keras.losses.SparseCategoricalCrossentropy(),
                 steps_per_execution=gradient_accumulation_steps * 2)
       m.set_gradient_accumulation_options(