docstr: Train this tensorflow model with train data.
  
  :param data: train data. It can be XShards, Spark DataFrame, Ray Dataset 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.
  :param steps_per_epoch: Total number of steps (batches of samples) before declaring one
         epoch finished and starting the next epoch. If `steps_pre_epoch` is `None`, the
         epoch will run until the input dataset is exhausted. When passing an infinitely
         repeating dataset, you must specify the `step_per_epoch` argument.
  :param validation_steps: Total number of steps (batches of samples) to draw before stopping
         when performing validation at the end of every epoch. Default: None.
  :param validation_freq: Only relevant if validation data is provided. Integer of
         `collections_abc.Container` instance (e.g. list, tuple, etc.). If an integer,
         specifies how many training epochs to run before a new validation run is performed,
         e.g. `validation_freq=2` runs validation every 2 epochs. If a Container, specifies
         the epochs on which to run validation, e.g. `validation_freq=[1, 2, 10]` runs
         validation at the end of the 1st, 2nd, and 10th epochs.
  :param data_config: An optional dictionary that can be passed to data creator function.
         If data is a Ray Dataset, specifies `output_signature` same as in
         `tf.data.Dataset.from_generator` (If `label_cols` is specified, a 2-element
         tuple of `tf.TypeSpec` objects corresponding to (features, label). Otherwise,
         a single `tf.TypeSpec` corresponding to features tensor).
  :param feature_cols: Feature column name(s) of data. Only used when data is a Spark
         DataFrame, an XShards of Pandas DataFrame or a Ray Dataset. Default: None.
  :param label_cols: Label column name(s) of data. Only used when data is a Spark DataFrame,
         an XShards of Pandas DataFrame or a Ray Dataset.
         Default: None.
  :return:
code: def fit(
      self,
      data: Union[
          "SparkXShards",
          "SparkDataFrame",
          "TFDataset",
          "ray.data.Dataset",
          Callable,
      ],
      epochs: int = 1,
      batch_size: int = 32,
      verbose: Union[str, int] = 1,
      callbacks: Optional[List["Callback"]] = None,
      validation_data: Optional[
          Union[
              "SparkXShards",
              "SparkDataFrame",
              "TFDataset",
              "ray.data.Dataset",
              Callable,
          ]
      ] = 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,
  ) -> Union[List[Dict[str, float]], Dict[str, float]]:
      """
      Train this tensorflow model with train data.
  
      :param data: train data. It can be XShards, Spark DataFrame, Ray Dataset 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.
      :param steps_per_epoch: Total number of steps (batches of samples) before declaring one
             epoch finished and starting the next epoch. If `steps_pre_epoch` is `None`, the
             epoch will run until the input dataset is exhausted. When passing an infinitely
             repeating dataset, you must specify the `step_per_epoch` argument.
      :param validation_steps: Total number of steps (batches of samples) to draw before stopping
             when performing validation at the end of every epoch. Default: None.
      :param validation_freq: Only relevant if validation data is provided. Integer of
             `collections_abc.Container` instance (e.g. list, tuple, etc.). If an integer,
             specifies how many training epochs to run before a new validation run is performed,
             e.g. `validation_freq=2` runs validation every 2 epochs. If a Container, specifies
             the epochs on which to run validation, e.g. `validation_freq=[1, 2, 10]` runs
             validation at the end of the 1st, 2nd, and 10th epochs.
      :param data_config: An optional dictionary that can be passed to data creator function.
             If data is a Ray Dataset, specifies `output_signature` same as in
             `tf.data.Dataset.from_generator` (If `label_cols` is specified, a 2-element
             tuple of `tf.TypeSpec` objects corresponding to (features, label). Otherwise,
             a single `tf.TypeSpec` corresponding to features tensor).
      :param feature_cols: Feature column name(s) of data. Only used when data is a Spark
             DataFrame, an XShards of Pandas DataFrame or a Ray Dataset. Default: None.
      :param label_cols: Label column name(s) of data. Only used when data is a Spark DataFrame,
             an XShards of Pandas DataFrame or a Ray Dataset.
             Default: None.
      :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
      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,
      )
  
      from bigdl.orca.data import SparkXShards
      from bigdl.orca.data.tf.data import Dataset
      from bigdl.orca.data.tf.tf2_data import TF2Dataset
  
      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):
          # Make sure each worker can get at least one data partition
          if data.num_partitions() < self.num_workers:
              data = data.repartition(self.num_workers)
          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"
              )
          ray_xshards = process_spark_xshards(
              data, self.num_workers
          )  # type:ignore
          val_ray_xshards = None
          if validation_data is not None:
              val_ray_xshards = process_spark_xshards(
                  validation_data,  # type:ignore
                  self.num_workers,
              )
  
          worker_stats = self._fit_ray_xshards(
              ray_xshards, val_ray_xshards, params
          )
      elif isinstance(data, Dataset):
          ray_xshards = TF2Dataset(data).get_ray_xshards(
              self.num_workers
          )
          val_ray_xshards = None
          if validation_data is not None:
              invalidInputError(
                  isinstance(validation_data, Dataset),
                  "Validation data type should be the same as train data,"
                  " but got type: {}".format(type(validation_data)),
              )
              val_ray_xshards = TF2Dataset(
                  validation_data
              ).get_ray_xshards(
                  self.num_workers
              )  # type:ignore
  
          worker_stats = self._fit_ray_xshards(
              ray_xshards, val_ray_xshards, params
          )
      elif isinstance(data, ray.data.Dataset):
          shards = data.split(
              n=self.num_workers, locality_hints=self.remote_workers
          )
  
          remote_worker_stats = []
          if validation_data is None:
              for shard, worker in zip(shards, self.remote_workers):
                  params["data_creator"] = self.process_ray_dataset(
                      shard, label_cols, feature_cols, data_config
                  )
                  remote_worker_stats.append(
                      worker.step.remote(**params)
                  )
              worker_stats = ray.get(remote_worker_stats)
          else:
              invalidInputError(
                  isinstance(validation_data, ray.data.Dataset),
                  "Validation data type should be the same as train data,"
                  " but got type: {}".format(type(validation_data)),
              )
  
              val_shards = validation_data.split(
                  n=self.num_workers,  # type:ignore
                  locality_hints=self.remote_workers,
              )
  
              for i in range(self.num_workers):
                  params["data_creator"] = self.process_ray_dataset(
                      shards[i], label_cols, feature_cols, data_config
                  )
                  params[
                      "validation_data_creator"
                  ] = self.process_ray_dataset(
                      val_shards[i],  # type:ignore
                      label_cols,
                      feature_cols,
                      data_config,
                  )
                  remote_worker_stats.append(
                      self.remote_workers[i].step.remote(**params)
                  )
              worker_stats = ray.get(remote_worker_stats)
      else:
          params["data_creator"] = data  # type:ignore
          params[
              "validation_data_creator"
          ] = validation_data  # type:ignore
          params_list = [params] * self.num_workers
  
          worker_stats = ray.get(
              [
                  self.remote_workers[i].step.remote(**params_list[i])
                  for i in range(self.num_workers)
              ]
          )
      # TensorFlow automatically synchronizes results on all the workers
      # and thus only need to return the result of the first worker
      stats = worker_stats[0].copy()
      return stats
  
imports: import os
  import types
  from bigdl.orca.common import SafePickle
  import shutil
  import tempfile
  import logging
  import ray
  from bigdl.dllib.utils import log4Error
  from bigdl.dllib.utils.file_utils import is_local_path
  from bigdl.orca.data.file import enable_multi_fs_save, enable_multi_fs_load
  from bigdl.orca.data.ray_xshards import RayXShards
  from bigdl.orca.learn.dl_cluster import RayDLCluster
  from bigdl.orca.learn.tf2.tf_runner import TFRunner
  from bigdl.orca.learn.ray_estimator import Estimator
  from bigdl.orca.learn.utils import maybe_dataframe_to_xshards, dataframe_to_xshards, convert_predict_xshards_to_dataframe, update_predict_xshards, process_xshards_of_pandas_dataframe, make_data_creator, add_predict_to_pd_xshards
  from bigdl.orca.data.file import get_remote_file_to_local, get_remote_dir_to_local, is_file
  from bigdl.orca.data.utils import process_spark_xshards
  from bigdl.dllib.utils.log4Error import invalidInputError
  from bigdl.orca.ray import OrcaRayContext
  from typing import TYPE_CHECKING, Dict, List, Callable, Union, Optional
label_desc: Use image node_id 2 for calling the TensorFlow2Estimator 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: stats