docstr: Train this tensorflow model with train data.
  :param data: train data. It can be XShards, Spark DataFrame or creator function which
         returns Iter or DataLoader.
         If data is XShards, each partition can be a Pandas DataFrame or a dictionary of
         {'x': feature, 'y': label}, where feature(label) is a numpy array or a tuple of
         numpy arrays.
  :param epochs: Number of epochs to train the model. Default: 1.
  :param batch_size: Total batch size for all workers used for training. Each worker's batch
         size would be this value divide the total number of workers. Default: 32.
  :param verbose: Prints output of one model if true.
  :param callbacks: List of Keras compatible callbacks to apply during training.
  :param validation_data: validation data. Validation data type should be the same
         as train data.
  :param class_weight: Optional dictionary mapping class indices (integers) to a weight
         (float) value, used for weighting the loss function. This can be useful to tell
         the model to "pay more attention" to samples from an under-represented class.
  :return:
docstr_tok: ["Train","this","tensorflow","model","with","train","data",".",":","param","data",":","train","data",".","It","can","be","XShards",",","Spark","DataFrame","or","creator","function","which","returns","Iter","or","DataLoader",".","If","data","is","XShards",",","each","partition","can","be","a","Pandas","DataFrame","or","a","dictionary","of","{","'","x","'",":","feature",",","'","y","'",":","label","}",",","where","feature","(","label",")","is","a","numpy","array","or","a","tuple","of","numpy","arrays",".",":","param","epochs",":","Number","of","epochs","to","train","the","model",".","Default",":","1.",":","param","batch_size",":","Total","batch","size","for","all","workers","used","for","training",".","Each","worker","'s","batch","size","would","be","this","value","divide","the","total","number","of","workers",".","Default",":","32.",":","param","verbose",":","Prints","output","of","one","model","if","true",".",":","param","callbacks",":","List","of","Keras","compatible","callbacks","to","apply","during","training",".",":","param","validation_data",":","validation","data",".","Validation","data","type","should","be","the","same","as","train","data",".",":","param","class_weight",":","Optional","dictionary","mapping","class","indices","(","integers",")","to","a","weight","(","float",")","value",",","used","for","weighting","the","loss","function",".","This","can","be","useful","to","tell","the","model","to","``","pay","more","attention","''","to","samples","from","an","under-represented","class",".",":","return",":"]
code: def fit(
      self,
      data: Union["SparkXShards", "SparkDataFrame", Callable],
      epochs: int = 1,
      batch_size: int = 32,
      verbose: Union[str, int] = 1,
      callbacks: Optional[List["Callback"]] = None,
      validation_data: Union[
          "SparkXShards", "SparkDataFrame", Callable, None
      ] = None,
      class_weight: Optional[Dict[int, float]] = None,
      initial_epoch: int = 0,
      steps_per_epoch: Optional[int] = None,
      validation_steps: Optional[int] = None,
      validation_freq: int = 1,
      data_config: Optional[Dict] = None,
      feature_cols: Optional[List[str]] = None,
      label_cols: Optional[List[str]] = None,
  ) -> Dict:
      """
      Train this tensorflow model with train data.
      :param data: train data. It can be XShards, Spark DataFrame or creator function which
             returns Iter or DataLoader.
             If data is XShards, each partition can be a Pandas DataFrame or a dictionary of
             {'x': feature, 'y': label}, where feature(label) is a numpy array or a tuple of
             numpy arrays.
      :param epochs: Number of epochs to train the model. Default: 1.
      :param batch_size: Total batch size for all workers used for training. Each worker's batch
             size would be this value divide the total number of workers. Default: 32.
      :param verbose: Prints output of one model if true.
      :param callbacks: List of Keras compatible callbacks to apply during training.
      :param validation_data: validation data. Validation data type should be the same
             as train data.
      :param class_weight: Optional dictionary mapping class indices (integers) to a weight
             (float) value, used for weighting the loss function. This can be useful to tell
             the model to "pay more attention" to samples from an under-represented class.
      :return:
      """
      if not isinstance(data, types.FunctionType):
          invalidInputError(
              isinstance(batch_size, int) and batch_size > 0,
              "batch_size should be a positive integer",
          )
      else:
          # batch_size can be None if the return of data_creator already generates batches
          if batch_size:
              invalidInputError(
                  isinstance(batch_size, int) and batch_size > 0,
                  "batch_size should be a positive integer",
              )
      # Use the local batch size for each worker to convert to XShards
      if batch_size:
          local_batch_size = batch_size // self.num_workers
          if local_batch_size <= 0:
              local_batch_size = 1
      else:
          local_batch_size = None
      sc = OrcaContext.get_spark_context()
  
      # for continuous training
      if self.model_weights:
          weights = sc.broadcast(self.model_weights)
      else:
          weights = None
  
      if self.optimizer_weights:
          opt_weights = sc.broadcast(self.optimizer_weights)
      else:
          opt_weights = None
  
      init_params = dict(
          model_creator=self.model_creator,
          model_load=self.load_path,
          compile_args_creator=self.compile_args_creator,
          config=self.config,
          verbose=self.verbose,
          size=self.num_workers,
          model_weights=weights,
          optimizer_weights=opt_weights,
          mode="fit",
          cluster_info=self._get_cluster_info(sc),
          model_dir=self.model_dir,
          application_id=self.application_id,
          need_to_log_to_driver=self.need_to_log_to_driver,
          driver_ip=self.ip,
          driver_port=self.port,
      )
  
      if self.load_params is not None:
          init_params.update(self.load_params)
  
      params = dict(
          epochs=epochs,
          batch_size=batch_size,
          verbose=verbose,
          callbacks=callbacks,
          class_weight=class_weight,
          initial_epoch=initial_epoch,
          steps_per_epoch=steps_per_epoch,
          validation_steps=validation_steps,
          validation_freq=validation_freq,
          data_config=data_config,
      )
  
      if isinstance(data, SparkXShards):
          data = data.to_lazy()
          if validation_data is not None and isinstance(
              validation_data, SparkXShards
          ):
              validation_data = validation_data.to_lazy()
      # Data partition should be equal to num workers.
      # Repartition Spark DataFrame before converting to SparkXShards.
      # Repartition on SparkXShards will result in empty partitions.
      if isinstance(data, DataFrame) or isinstance(data, SparkXShards):
          if data.rdd.getNumPartitions() != self.num_workers:
              data = data.repartition(self.num_workers)
          if validation_data is not None:
              invalidInputError(
                  isinstance(validation_data, DataFrame)
                  or isinstance(validation_data, SparkXShards),
                  "validation_data should have the same type with train data",
              )
              if (
                  validation_data.rdd.getNumPartitions()
                  != self.num_workers
              ):  # type:ignore
                  validation_data = validation_data.repartition(
                      self.num_workers
                  )  # type:ignore
      data, validation_data = maybe_dataframe_to_xshards(
          data,
          validation_data,
          feature_cols,
          label_cols,
          mode="fit",
          num_workers=self.num_workers,
          accept_str_col=True,
          shard_size=local_batch_size,
      )
  
      if isinstance(
          data, SparkXShards
      ):  # Computation triggered when collect
          if data._get_class_name() == "pandas.core.frame.DataFrame":
              (
                  data,
                  validation_data,
              ) = process_xshards_of_pandas_dataframe(
                  data, feature_cols, label_cols, validation_data, "fit"
              )
          if validation_data is None:
  
              def transform_func(iter, init_param, param):
                  partition_data = list(iter)
                  param["data_creator"] = make_data_creator(
                      partition_data
                  )
                  return SparkRunner(**init_param).step(**param)
  
              res = (
                  data.rdd.barrier()
                  .mapPartitions(  # type:ignore
                      lambda iter: transform_func(
                          iter, init_params, params
                      )
                  )
                  .collect()
              )
          else:
  
              def transform_func(iter, init_param, param):
                  data_tuple_list = list(iter)
                  data_list = [
                      x
                      for data_tuple in data_tuple_list
                      for x in data_tuple[0]
                  ]
                  valid_list = [
                      x
                      for data_tuple in data_tuple_list
                      for x in data_tuple[1]
                  ]
                  param["data_creator"] = make_data_creator(data_list)
                  param["validation_data_creator"] = make_data_creator(
                      valid_list
                  )
                  return SparkRunner(**init_param).step(**param)
  
              train_rdd = data.rdd.mapPartitions(
                  lambda iter: [list(iter)]
              )  # type:ignore
              val_rdd = (
                  validation_data.rdd.mapPartitions(  # type:ignore
                      lambda iter: [list(iter)]
                  )
              )
              res = (
                  train_rdd.zip(val_rdd)
                  .barrier()
                  .mapPartitions(
                      lambda iter: transform_func(
                          iter, init_params, params
                      )
                  )
                  .collect()
              )
      else:
          params["data_creator"] = data  # type:ignore
          params[
              "validation_data_creator"
          ] = validation_data  # type:ignore
  
          def transform_func(iter, init_param, param):
              return SparkRunner(**init_param).step(**param)
  
          res = (
              self.workerRDD.barrier()
              .mapPartitions(
                  lambda iter: transform_func(iter, init_params, params)
              )
              .collect()
          )
  
      result = self._update_weights(res)
  
      return result
  
imports: import os
  import types
  import logging
  import tempfile
  import shutil
  from bigdl.dllib.utils.common import get_node_and_core_number
  from bigdl.dllib.utils.file_utils import is_local_path
  from bigdl.dllib.utils.utils import get_node_ip
  from bigdl.orca.data.file import is_file, exists, get_remote_file_to_local, get_remote_files_with_prefix_to_local, put_local_file_to_remote, put_local_files_with_prefix_to_remote
  from bigdl.orca.learn.tf2.spark_runner import SparkRunner
  from bigdl.orca.learn.utils import find_free_port, find_ip_and_free_port
  from bigdl.orca.learn.utils import maybe_dataframe_to_xshards, dataframe_to_xshards, convert_predict_xshards_to_dataframe, make_data_creator, load_model, save_model, process_xshards_of_pandas_dataframe, add_predict_to_pd_xshards, update_predict_xshards
  from bigdl.orca.learn.log_monitor import start_log_server, stop_log_server
  from bigdl.orca import OrcaContext
  from bigdl.orca.data.shard import SparkXShards
  from bigdl.dllib.utils.log4Error import invalidInputError
  from pyspark.sql.dataframe import DataFrame
  from typing import TYPE_CHECKING, Any, Dict, List, Callable, Union, Optional
label_desc: Use image node_id 3 for calling the SparkTFEstimator obj's underlying member method code with example usage: obj.fit(data, epochs, batch_size, verbose, callbacks, validation_data, class_weight, initial_epoch, steps_per_epoch, validation_steps, validation_freq, data_config, feature_cols, label_cols) and returns: result, SparkRunner, SparkRunner, SparkRunner