Supporting models with required training arg in call

Summary:
Support keras models with call signatures which have no default value for the `training` parameter.
These are supported in upstream keras models by checking the model's call signature and passing False if the argument is required but has not been passed.
TF2.5 Only

Test Plan:
Added a test which checks that you can call build/fit/evaluate/predict on models with a required training parameter.
Added a test which checks that you get an appropriate error message when call requires other extra arguments.

Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, georgep

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

Maniphest Tasks: T57217

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

Author: Alfie-Edwards

tensorflow/python/ipu/keras/extensions/model_extensions.py

@@ -215,6 +215,7 @@ def to_input_tensor(x):
           return x
         return input_layer.Input(batch_shape=x)
 
+      self._check_call_args('build')
       inputs = map_shapes(input_shape, to_input_tensor)
       outputs = self._trace_graph_network(inputs)
       self._update_graph_network(inputs, outputs)
@@ -547,6 +548,18 @@ def _trace_graph_network(self, inputs):
       layer._inbound_nodes.clear()  # pylint: disable=protected-access
       layer._outbound_nodes.clear()  # pylint: disable=protected-access
 
+    # Case where `training` is a positional arg with no default.
+    kwargs = {}
+    # Update self._call_full_argspec.
+    self._init_call_fn_args()
+    call_signature = self._call_full_argspec
+    if len(call_signature.args) > 2:
+      n_required = len(call_signature.args) - len(call_signature.defaults
+                                                  or [])
+      for arg in call_signature.args[2:n_required]:
+        if arg == 'training':
+          kwargs['training'] = False
+
     # Invoke call within a call_context in case metrics are added in call.
     # We can't invoke self() because the base_layer __call__ override sees that
     # we are passing in keras tensors and treats the whole model as a layer
@@ -567,7 +580,7 @@ def get_shape(tensor):
           input_shape = get_shape(inputs)
 
         self.build(input_shape)
-      return self.call(inputs)
+      return self.call(inputs, **kwargs)
 
   @trackable.no_automatic_dependency_tracking
   def _update_graph_network(self, inputs, outputs):

tensorflow/python/ipu/tests/keras/extensions/single_ipu/functional_test.py

@@ -570,6 +570,40 @@ def call(self, x):  # pylint: disable=arguments-differ
     output = model(sample_input)  # pylint: disable=not-callable
     self.assertEqual(output.shape, (1, 3))
 
+  def testModelWithRequiredTrainingArg(self):
+    class MyModel(training_lib.Model):  # pylint: disable=abstract-method
+      def call(self, inputs, training):  # pylint: disable=arguments-differ,unused-argument
+        return inputs
+
+    model = MyModel()
+    model.compile()
+    x = np.ones(8)
+    y = np.ones(8)
+
+    model.build([8])
+    model.fit(x, y, batch_size=1)
+    model.evaluate(x, y, batch_size=1)
+    model.predict(x, batch_size=1)
+
+  def testModelWithAdditionalRequiredArgs(self):
+    class MyModel(training_lib.Model):  # pylint: disable=abstract-method
+      def call(self, inputs, additional_arg):  # pylint: disable=arguments-differ,unused-argument
+        return inputs
+
+    model = MyModel()
+    model.compile()
+    x = np.ones(8)
+    y = np.ones(8)
+
+    with self.assertRaisesRegex(ValueError, r'Models passed to `build`'):
+      model.build([8])
+    with self.assertRaisesRegex(ValueError, r'Models passed to `fit`'):
+      model.fit(x, y, batch_size=1)
+    with self.assertRaisesRegex(ValueError, r'Models passed to `evaluate`'):
+      model.evaluate(x, y, batch_size=1)
+    with self.assertRaisesRegex(ValueError, r'Models passed to `predict`'):
+      model.predict(x, batch_size=1)
+
   def testNoneInShapeWithFunctionalAPI(self):
     # pylint: disable=abstract-method
     class BasicBlock(training_lib.Model):