class_docstr: Base context class that contains runtime information for any Determined
  workflow that uses the ``tf.keras`` API.
class_docstr_tok: ["Base","context","class","that","contains","runtime","information","for","any","Determined","workflow","that","uses","the","``","tf.keras","``","API","."]
code: def _wrap_model_with_train_fn(
      self, model: Any, train_fn: Optional[Callable]
  ) -> Any:
      class _WrappedModel(type(model)):  # type: ignore
          def __init__(wrapper) -> None:
              self.model = model
  
          def __getattr__(wrapper, name):  # type: ignore
              return getattr(model, name)
  
          def __setattr__(wrapper, name, value):  # type: ignore
              return setattr(model, name, value)
  
          def __delattr__(wrapper, name):  # type: ignore
              return delattr(model, name)
  
          def compile(wrapper, *args: Any, **kwargs: Any) -> None:
              bound_arguments = inspect.signature(model.compile).bind(
                  *args, **kwargs
              )
              bound_arguments.apply_defaults()
              self.compile_args = bound_arguments
  
          def fit_generator(wrapper, *args: Any, **kwargs: Any) -> None:
              if not self.compile_args:
                  raise errors.InvalidExperimentException(
                      "Must call .compile before calling .fit_generator()."
                  )
  
              fit_generator_args = inspect.signature(
                  model.fit_generator
              ).bind(*args, **kwargs)
              fit_generator_args.apply_defaults()
  
              training_data = fit_generator_args.arguments["generator"]
  
              if (
                  fit_generator_args.arguments["validation_data"]
                  is None
              ):
                  raise errors.InvalidExperimentException(
                      "Determined requires validation_data in the call to fit_generator()."
                  )
  
              validation_data = keras._adapt_data_from_data_loader(
                  input_data=fit_generator_args.arguments[
                      "validation_data"
                  ],
                  batch_size=self._per_slot_batch_size,
              )
  
              self.train_config = TFKerasTrainConfig(
                  training_data=training_data,
                  validation_data=validation_data,
                  callbacks=fit_generator_args.arguments["callbacks"],
              )
  
              self.configure_fit(
                  verbose=fit_generator_args.arguments["verbose"],
                  class_weight=fit_generator_args.arguments[
                      "class_weight"
                  ],
                  shuffle=fit_generator_args.arguments["shuffle"],
                  workers=fit_generator_args.arguments["workers"],
                  use_multiprocessing=fit_generator_args.arguments[
                      "use_multiprocessing"
                  ],
                  max_queue_size=fit_generator_args.arguments[
                      "max_queue_size"
                  ],
              )
  
              if train_fn:
                  train_fn()
  
          def fit(wrapper, *args: Any, **kwargs: Any) -> None:
              """Communicate a model, data, and other training configuration with the harness.
  
              Parameters:
                  the same as tf.keras.Model.fit except for this function only handles the
                  following cases of data
  
                  x: Input data. It could be:
                      1) A Numpy array (or array-like), or a list of arrays (in case the model
                      has multiple inputs).
                      2) A dict mapping input names to the corresponding array, if the model
                      has named inputs.
                      3) A tf.data dataset. Should return a tuple of either (inputs, targets) or
                      (inputs, targets, sample_weights).
                      4) A keras.utils.Sequence returning (inputs, targets) or (inputs, targets,
                      sample weights).
  
                  y: Target data. Like the input data x, it could be either Numpy array(s).
                      If x is a dataset or keras.utils.Sequence instance, y should not be
                      specified(since targets will be obtained from x).
  
                  validation_data: Data on which to evaluate the loss and any model metrics
                      at the end of each epoch. The model will not be trained on this data.
                      validation_data will override validation_split. validation_data could be:
                      1) tuple (x_val, y_val) of Numpy arrays
                      2) tuple (x_val, y_val, val_sample_weights) of Numpy arrays
                      3) dataset For the first two cases, batch_size must be provided.
                      For the last case, validation_steps could be provided.
              """
              if not self.compile_args:
                  raise errors.InvalidExperimentException(
                      "Must call .compile before calling .fit()."
                  )
  
              fit_args = inspect.signature(model.fit).bind(
                  *args, **kwargs
              )
              fit_args.apply_defaults()
  
              training_data = keras._adapt_data_from_fit_args(
                  x=fit_args.arguments["x"],
                  y=fit_args.arguments["y"],
                  sample_weight=fit_args.arguments["sample_weight"],
                  batch_size=self._per_slot_batch_size,
              )
  
              if fit_args.arguments["validation_data"] is None:
                  raise errors.InvalidExperimentException(
                      "Determined requires validation_data in the call to fit()."
                  )
  
              validation_data = keras._adapt_data_from_data_loader(
                  input_data=fit_args.arguments["validation_data"],
                  batch_size=self._per_slot_batch_size,
              )
  
              self.train_config = TFKerasTrainConfig(
                  training_data=training_data,
                  validation_data=validation_data,
                  callbacks=fit_args.arguments["callbacks"],
              )
  
              self.configure_fit(
                  verbose=fit_args.arguments["verbose"],
                  shuffle=fit_args.arguments["shuffle"],
                  class_weight=fit_args.arguments["class_weight"],
                  workers=fit_args.arguments["workers"],
                  use_multiprocessing=fit_args.arguments[
                      "use_multiprocessing"
                  ],
                  max_queue_size=fit_args.arguments["max_queue_size"],
              )
  
              if train_fn:
                  train_fn()
  
      return _WrappedModel()
  
imports: import inspect
  import logging
  from typing import Any, Callable, Dict, List, NamedTuple, Optional, Union, cast
  import tensorflow
  import determined
  from determined import errors, keras, util
  from determined.common import check
  from determined.horovod import hvd
label_desc: Use image node_id 5 for calling the TFKerasTrialContext obj's underlying member method code with example usage: obj._wrap_model_with_train_fn(model, train_fn) and returns: _WrappedModel, getattr, setattr, delattr