class_docstr: Groups wrapped layer call functions.
  
  This is used to ensure that all layer call functions are traced with the same
  inputs-
    - call
    - call_and_return_conditional_losses
    - call_and_return_all_conditional_losses
class_docstr_tok: ["Groups","wrapped","layer","call","functions",".","This","is","used","to","ensure","that","all","layer","call","functions","are","traced","with","the","same","inputs-","-","call","-","call_and_return_conditional_losses","-","call_and_return_all_conditional_losses"]
docstr: Wraps call function with added training argument if necessary.
docstr_tok: ["Wraps","call","function","with","added","training","argument","if","necessary","."]
code: def _maybe_wrap_with_training_arg(
      self, call_fn, match_layer_training_arg
  ):
      """Wraps call function with added training argument if necessary."""
      if (
          not self.layer._expects_training_arg
          and self._expects_training_arg
      ):  # pylint: disable=protected-access
          # Add training arg to wrapper function.
          arg_spec = tf_inspect.getfullargspec(call_fn)
          args = arg_spec.args + ["training"]
          defaults = list(arg_spec.defaults or [])
          defaults.append(False)
          new_arg_spec = tf_inspect.FullArgSpec(
              args=args,
              varargs=arg_spec.varargs,
              varkw=arg_spec.varkw,
              defaults=defaults,
              kwonlyargs=arg_spec.kwonlyargs,
              kwonlydefaults=arg_spec.kwonlydefaults,
              annotations=arg_spec.annotations,
          )
  
          # Set new training arg index
          self._training_arg_index = len(args) - 1
          if tf_inspect.ismethod(call_fn):
              self._training_arg_index -= 1
  
          def wrap_with_training_arg(*args, **kwargs):
              if match_layer_training_arg:
                  # Remove the training value, since the original call_fn does not
                  # expect a training arg. Instead, the training value will be
                  # propagated using the call context created in LayerCall.
                  args = list(args)
                  kwargs = kwargs.copy()
                  utils.remove_training_arg(
                      self._training_arg_index, args, kwargs
                  )
              return call_fn(*args, **kwargs)
  
          return tf_decorator.make_decorator(
              target=call_fn,
              decorator_func=wrap_with_training_arg,
              decorator_argspec=new_arg_spec,
          )
  
      return call_fn
  
imports: import functools
  import threading
  import weakref
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.keras import backend
  from tensorflow.python.keras.engine import base_layer_utils
  from tensorflow.python.keras.engine import input_spec
  from tensorflow.python.keras.mixed_precision import autocast_variable
  from tensorflow.python.keras.saving import saving_utils
  from tensorflow.python.keras.saving.saved_model import constants
  from tensorflow.python.keras.saving.saved_model import load
  from tensorflow.python.keras.saving.saved_model import serialized_attributes
  from tensorflow.python.keras.saving.saved_model import utils
  from tensorflow.python.keras.utils import tf_contextlib
  from tensorflow.python.keras.utils import tf_inspect
  from tensorflow.python.keras.utils import tf_utils
  from tensorflow.python.keras.utils import version_utils
  from tensorflow.python.keras.utils.generic_utils import LazyLoader
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.trackable import data_structures
  from tensorflow.python.util import nest
  from tensorflow.python.util import tf_decorator
label_desc: Use image node_id 8 for calling the LayerCallCollection obj's underlying member method code with example usage: obj._maybe_wrap_with_training_arg(call_fn, match_layer_training_arg) and returns: call_fn, tf_decorator, call_fn