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","."]
docstr: This should be used to wrap ``tf.data.Dataset`` objects immediately after
  they have been created. Users should use the output of this wrapper as the
  new instance of their dataset. If users create multiple datasets (e.g.,
  one for training and one for validation), users should wrap each dataset
  independently.
  
  Args:
      dataset: tf.data.Dataset
      shard_dataset:
          When performing multi-slot (distributed) training, this
          controls whether the dataset is sharded so that each training process
          (one per slot) sees unique data. If set to False, users must manually
          configure each process to use unique data.
docstr_tok: ["This","should","be","used","to","wrap","``","tf.data.Dataset","``","objects","immediately","after","they","have","been","created",".","Users","should","use","the","output","of","this","wrapper","as","the","new","instance","of","their","dataset",".","If","users","create","multiple","datasets","(","e.g.",",","one","for","training","and","one","for","validation",")",",","users","should","wrap","each","dataset","independently",".","Args",":","dataset",":","tf.data.Dataset","shard_dataset",":","When","performing","multi-slot","(","distributed",")","training",",","this","controls","whether","the","dataset","is","sharded","so","that","each","training","process","(","one","per","slot",")","sees","unique","data",".","If","set","to","False",",","users","must","manually","configure","each","process","to","use","unique","data","."]
code: def wrap_dataset(
      self, dataset: Any, shard_dataset: bool = True
  ) -> Any:
      """
      This should be used to wrap ``tf.data.Dataset`` objects immediately after
      they have been created. Users should use the output of this wrapper as the
      new instance of their dataset. If users create multiple datasets (e.g.,
      one for training and one for validation), users should wrap each dataset
      independently.
  
      Args:
          dataset: tf.data.Dataset
          shard_dataset:
              When performing multi-slot (distributed) training, this
              controls whether the dataset is sharded so that each training process
              (one per slot) sees unique data. If set to False, users must manually
              configure each process to use unique data.
      """
      if not self.env.managed_training:
          return dataset
  
      self.dataset_initialized = True
      if (
          self.distributed.size == 1
          or not isinstance(dataset, tf.data.Dataset)
          or not shard_dataset
      ):
          if self.distributed.size > 1 and not shard_dataset:
              logger.info("Dataset sharding skipped.")
          return dataset
  
      hvd.require_horovod_type(
          "tensorflow.keras",
          "TFKerasTrialContext.wrap_dataset was called.",
      )
      dataset = dataset.shard(hvd.size(), hvd.rank())
      logger.debug(
          f"Sharded dataset to index {hvd.rank()} of {hvd.size()}."
      )
      return dataset
  
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 6 for calling the TFKerasTrialContext obj's underlying member method code with example usage: obj.wrap_dataset(dataset, shard_dataset) and returns: dataset, dataset, dataset