class_docstr: An iterator for reading records from a :class:`~Dataset`.
  
  For Datasets, each iteration call represents a complete read of all items in the
  Dataset.
  
  If using Ray Train, each trainer actor should get its own iterator by calling
  :meth:`ray.train.get_dataset_shard("train")
  <ray.train.get_dataset_shard>`.
  
  Examples:
      >>> import ray
      >>> ds = ray.data.range(5)
      >>> ds
      Dataset(num_blocks=..., num_rows=5, schema={id: int64})
      >>> ds.iterator()
      DataIterator(Dataset(num_blocks=..., num_rows=5, schema={id: int64}))
class_docstr_tok: ["An","iterator","for","reading","records","from","a",":","class",":","`","~Dataset","`",".","For","Datasets",",","each","iteration","call","represents","a","complete","read","of","all","items","in","the","Dataset",".","If","using","Ray","Train",",","each","trainer","actor","should","get","its","own","iterator","by","calling",":","meth",":","`","ray.train.get_dataset_shard","(","``","train","''",")","<","ray.train.get_dataset_shard",">","`",".","Examples",":",">",">",">","import","ray",">",">",">","ds","=","ray.data.range","(","5",")",">",">",">","ds","Dataset","(","num_blocks=","...",",","num_rows=5",",","schema=","{","id",":","int64","}",")",">",">",">","ds.iterator","(",")","DataIterator","(","Dataset","(","num_blocks=","...",",","num_rows=5",",","schema=","{","id",":","int64","}",")",")"]
docstr: Return a batched iterable of Torch Tensors over the dataset.
  
  This iterable yields a dictionary of column-tensors. If you are looking for
  more flexibility in the tensor conversion (e.g. casting dtypes) or the batch
  format, try using :meth:`~ray.data.iterator.DataIterator.iter_batches` directly.
  
  Examples:
      >>> import ray
      >>> for batch in ray.data.range(
      ...     12,
      ... ).iterator().iter_torch_batches(batch_size=4):
      ...     print(batch)
      {'id': tensor([0, 1, 2, 3])}
      {'id': tensor([4, 5, 6, 7])}
      {'id': tensor([ 8,  9, 10, 11])}
  
      Use the ``collate_fn`` to customize how the tensor batch is created.
  
      >>> from typing import Any, Dict
      >>> import torch
      >>> import numpy as np
      >>> import ray
      >>> def collate_fn(batch: Dict[str, np.ndarray]) -> Any:
      ...     return torch.stack(
      ...         [torch.as_tensor(array) for array in batch.values()],
      ...         axis=1
      ...     )
      >>> iterator = ray.data.from_items([
      ...     {"col_1": 1, "col_2": 2},
      ...     {"col_1": 3, "col_2": 4}]).iterator()
      >>> for batch in iterator.iter_torch_batches(collate_fn=collate_fn):
      ...     print(batch)
      tensor([[1, 2],
              [3, 4]])
  
  Time complexity: O(1)
  
  Args:
      prefetch_batches: The number of batches to fetch ahead of the current batch
          to fetch. If set to greater than 0, a separate threadpool will be used
          to fetch the objects to the local node, format the batches, and apply
          the collate_fn. Defaults to 1.
      batch_size: The number of rows in each batch, or None to use entire blocks
          as batches (blocks may contain different number of rows).
          The final batch may include fewer than ``batch_size`` rows if
          ``drop_last`` is ``False``. Defaults to 256.
      dtypes: The Torch dtype(s) for the created tensor(s); if None, the dtype
          will be inferred from the tensor data. You can't use this parameter
          with ``collate_fn``.
      device: The device on which the tensor should be placed. Defaults to
          "auto" which moves the tensors to the appropriate device when the
          Dataset is passed to Ray Train and ``collate_fn`` is not provided.
          Otherwise, defaults to CPU. You can't use this parameter with
          ``collate_fn``.
      collate_fn: A function to convert a Numpy batch to a PyTorch tensor batch.
          When this parameter is specified, the user should manually handle the
          host to device data transfer outside of ``collate_fn``.
          This is useful for further processing the data after it has been
          batched. Potential use cases include collating along a dimension other
          than the first, padding sequences of various lengths, or generally
          handling batches of different length tensors. If not provided, the
          default collate function is used which simply converts the batch of
          numpy arrays to a batch of PyTorch tensors. This API is still
          experimental and is subject to change. You can't use this parameter in
          conjunction with ``dtypes`` or ``device``.
      drop_last: Whether to drop the last batch if it's incomplete.
      local_shuffle_buffer_size: If non-None, the data will be randomly shuffled
          using a local in-memory shuffle buffer, and this value will serve as the
          minimum number of rows that must be in the local in-memory shuffle
          buffer in order to yield a batch. When there are no more rows to add to
          the buffer, the remaining rows in the buffer will be drained. This
          buffer size must be greater than or equal to ``batch_size``, and
          therefore ``batch_size`` must also be specified when using local
          shuffling.
      local_shuffle_seed: The seed to use for the local random shuffle.
  
  Returns:
      An iterable over Torch Tensor batches.
docstr_tok: ["Return","a","batched","iterable","of","Torch","Tensors","over","the","dataset",".","This","iterable","yields","a","dictionary","of","column-tensors",".","If","you","are","looking","for","more","flexibility","in","the","tensor","conversion","(","e.g",".","casting","dtypes",")","or","the","batch","format",",","try","using",":","meth",":","`","~ray.data.iterator.DataIterator.iter_batches","`","directly",".","Examples",":",">",">",">","import","ray",">",">",">","for","batch","in","ray.data.range","(","...","12",",","...",")",".iterator","(",")",".iter_torch_batches","(","batch_size=4",")",":","...","print","(","batch",")","{","'id","'",":","tensor","(","[","0",",","1",",","2",",","3","]",")","}","{","'id","'",":","tensor","(","[","4",",","5",",","6",",","7","]",")","}","{","'id","'",":","tensor","(","[","8",",","9",",","10",",","11","]",")","}","Use","the","``","collate_fn","``","to","customize","how","the","tensor","batch","is","created",".",">",">",">","from","typing","import","Any",",","Dict",">",">",">","import","torch",">",">",">","import","numpy","as","np",">",">",">","import","ray",">",">",">","def","collate_fn","(","batch",":","Dict","[","str",",","np.ndarray","]",")","-",">","Any",":","...","return","torch.stack","(","...","[","torch.as_tensor","(","array",")","for","array","in","batch.values","(",")","]",",","...","axis=1","...",")",">",">",">","iterator","=","ray.data.from_items","(","[","...","{","``","col_1","''",":","1",",","``","col_2","''",":","2","}",",","...","{","``","col_1","''",":","3",",","``","col_2","''",":","4","}","]",")",".iterator","(",")",">",">",">","for","batch","in","iterator.iter_torch_batches","(","collate_fn=collate_fn",")",":","...","print","(","batch",")","tensor","(","[","[","1",",","2","]",",","[","3",",","4","]","]",")","Time","complexity",":","O","(","1",")","Args",":","prefetch_batches",":","The","number","of","batches","to","fetch","ahead","of","the","current","batch","to","fetch",".","If","set","to","greater","than","0",",","a","separate","threadpool","will","be","used","to","fetch","the","objects","to","the","local","node",",","format","the","batches",",","and","apply","the","collate_fn",".","Defaults","to","1.","batch_size",":","The","number","of","rows","in","each","batch",",","or","None","to","use","entire","blocks","as","batches","(","blocks","may","contain","different","number","of","rows",")",".","The","final","batch","may","include","fewer","than","``","batch_size","``","rows","if","``","drop_last","``","is","``","False","``",".","Defaults","to","256.","dtypes",":","The","Torch","dtype","(","s",")","for","the","created","tensor","(","s",")",";","if","None",",","the","dtype","will","be","inferred","from","the","tensor","data",".","You","ca","n't","use","this","parameter","with","``","collate_fn","``",".","device",":","The","device","on","which","the","tensor","should","be","placed",".","Defaults","to","``","auto","''","which","moves","the","tensors","to","the","appropriate","device","when","the","Dataset","is","passed","to","Ray","Train","and","``","collate_fn","``","is","not","provided",".","Otherwise",",","defaults","to","CPU",".","You","ca","n't","use","this","parameter","with","``","collate_fn","``",".","collate_fn",":","A","function","to","convert","a","Numpy","batch","to","a","PyTorch","tensor","batch",".","When","this","parameter","is","specified",",","the","user","should","manually","handle","the","host","to","device","data","transfer","outside","of","``","collate_fn","``",".","This","is","useful","for","further","processing","the","data","after","it","has","been","batched",".","Potential","use","cases","include","collating","along","a","dimension","other","than","the","first",",","padding","sequences","of","various","lengths",",","or","generally","handling","batches","of","different","length","tensors",".","If","not","provided",",","the","default","collate","function","is","used","which","simply","converts","the","batch","of","numpy","arrays","to","a","batch","of","PyTorch","tensors",".","This","API","is","still","experimental","and","is","subject","to","change",".","You","ca","n't","use","this","parameter","in","conjunction","with","``","dtypes","``","or","``","device","``",".","drop_last",":","Whether","to","drop","the","last","batch","if","it","'s","incomplete",".","local_shuffle_buffer_size",":","If","non-None",",","the","data","will","be","randomly","shuffled","using","a","local","in-memory","shuffle","buffer",",","and","this","value","will","serve","as","the","minimum","number","of","rows","that","must","be","in","the","local","in-memory","shuffle","buffer","in","order","to","yield","a","batch",".","When","there","are","no","more","rows","to","add","to","the","buffer",",","the","remaining","rows","in","the","buffer","will","be","drained",".","This","buffer","size","must","be","greater","than","or","equal","to","``","batch_size","``",",","and","therefore","``","batch_size","``","must","also","be","specified","when","using","local","shuffling",".","local_shuffle_seed",":","The","seed","to","use","for","the","local","random","shuffle",".","Returns",":","An","iterable","over","Torch","Tensor","batches","."]
code: def iter_torch_batches(
      self,
      *,
      prefetch_batches: int = 1,
      batch_size: Optional[int] = 256,
      dtypes: Optional[
          Union["torch.dtype", Dict[str, "torch.dtype"]]
      ] = None,
      device: str = "auto",
      collate_fn: Optional[
          Callable[[Dict[str, np.ndarray]], "CollatedData"]
      ] = None,
      drop_last: bool = False,
      local_shuffle_buffer_size: Optional[int] = None,
      local_shuffle_seed: Optional[int] = None,
      # Deprecated.
      prefetch_blocks: int = 0,
  ) -> Iterable["TorchBatchType"]:
      """Return a batched iterable of Torch Tensors over the dataset.
  
      This iterable yields a dictionary of column-tensors. If you are looking for
      more flexibility in the tensor conversion (e.g. casting dtypes) or the batch
      format, try using :meth:`~ray.data.iterator.DataIterator.iter_batches` directly.
  
      Examples:
          >>> import ray
          >>> for batch in ray.data.range(
          ...     12,
          ... ).iterator().iter_torch_batches(batch_size=4):
          ...     print(batch)
          {'id': tensor([0, 1, 2, 3])}
          {'id': tensor([4, 5, 6, 7])}
          {'id': tensor([ 8,  9, 10, 11])}
  
          Use the ``collate_fn`` to customize how the tensor batch is created.
  
          >>> from typing import Any, Dict
          >>> import torch
          >>> import numpy as np
          >>> import ray
          >>> def collate_fn(batch: Dict[str, np.ndarray]) -> Any:
          ...     return torch.stack(
          ...         [torch.as_tensor(array) for array in batch.values()],
          ...         axis=1
          ...     )
          >>> iterator = ray.data.from_items([
          ...     {"col_1": 1, "col_2": 2},
          ...     {"col_1": 3, "col_2": 4}]).iterator()
          >>> for batch in iterator.iter_torch_batches(collate_fn=collate_fn):
          ...     print(batch)
          tensor([[1, 2],
                  [3, 4]])
  
      Time complexity: O(1)
  
      Args:
          prefetch_batches: The number of batches to fetch ahead of the current batch
              to fetch. If set to greater than 0, a separate threadpool will be used
              to fetch the objects to the local node, format the batches, and apply
              the collate_fn. Defaults to 1.
          batch_size: The number of rows in each batch, or None to use entire blocks
              as batches (blocks may contain different number of rows).
              The final batch may include fewer than ``batch_size`` rows if
              ``drop_last`` is ``False``. Defaults to 256.
          dtypes: The Torch dtype(s) for the created tensor(s); if None, the dtype
              will be inferred from the tensor data. You can't use this parameter
              with ``collate_fn``.
          device: The device on which the tensor should be placed. Defaults to
              "auto" which moves the tensors to the appropriate device when the
              Dataset is passed to Ray Train and ``collate_fn`` is not provided.
              Otherwise, defaults to CPU. You can't use this parameter with
              ``collate_fn``.
          collate_fn: A function to convert a Numpy batch to a PyTorch tensor batch.
              When this parameter is specified, the user should manually handle the
              host to device data transfer outside of ``collate_fn``.
              This is useful for further processing the data after it has been
              batched. Potential use cases include collating along a dimension other
              than the first, padding sequences of various lengths, or generally
              handling batches of different length tensors. If not provided, the
              default collate function is used which simply converts the batch of
              numpy arrays to a batch of PyTorch tensors. This API is still
              experimental and is subject to change. You can't use this parameter in
              conjunction with ``dtypes`` or ``device``.
          drop_last: Whether to drop the last batch if it's incomplete.
          local_shuffle_buffer_size: If non-None, the data will be randomly shuffled
              using a local in-memory shuffle buffer, and this value will serve as the
              minimum number of rows that must be in the local in-memory shuffle
              buffer in order to yield a batch. When there are no more rows to add to
              the buffer, the remaining rows in the buffer will be drained. This
              buffer size must be greater than or equal to ``batch_size``, and
              therefore ``batch_size`` must also be specified when using local
              shuffling.
          local_shuffle_seed: The seed to use for the local random shuffle.
  
      Returns:
          An iterable over Torch Tensor batches.
      """
  
      from ray.air._internal.torch_utils import (
          convert_ndarray_batch_to_torch_tensor_batch,
          get_device,
      )
  
      if collate_fn is not None and (
          dtypes is not None or device != "auto"
      ):
          raise ValueError(
              "collate_fn cannot be used with dtypes and device."
              "You should manually move the output Torch tensors to the"
              "desired dtype and device outside of collate_fn."
          )
  
      if device == "auto":
          # Use the appropriate device for Ray Train, or falls back to CPU if
          # Ray Train is not being used.
          device = get_device()
  
      if collate_fn is None:
          # The default collate_fn handles formatting and Tensor creation.
          # Here, we set device=None to defer host to device data transfer
          # to the subsequent finalize_fn.
          def collate_fn(
              batch: Union[np.ndarray, Dict[str, np.ndarray]]
          ):
              return convert_ndarray_batch_to_torch_tensor_batch(
                  batch,
                  dtypes=dtypes,
                  device=None,
              )
  
          # The default finalize_fn handles the host to device data transfer.
          # This is executed in a 1-thread pool separately from collate_fn
          # to allow independent parallelism of these steps.
          def finalize_fn(
              batch: Union["torch.Tensor", Dict[str, "torch.Tensor"]]
          ):
              if device is not None:
                  if isinstance(batch, dict):
                      for k, t in batch.items():
                          batch[k] = t.to(device=device)
                  else:
                      batch = batch.to(device=device)
              return batch
  
      else:
          finalize_fn = None
  
      return self.iter_batches(
          prefetch_batches=prefetch_batches,
          prefetch_blocks=prefetch_blocks,
          batch_size=batch_size,
          drop_last=drop_last,
          local_shuffle_buffer_size=local_shuffle_buffer_size,
          local_shuffle_seed=local_shuffle_seed,
          _collate_fn=collate_fn,
          _finalize_fn=finalize_fn,
      )
  
imports: import abc
  import time
  from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple, TypeVar, Union
  import numpy
  from ray.data._internal.block_batching.iter_batches import iter_batches
  from ray.data._internal.stats import DatasetStats, StatsManager
  from ray.data.block import Block, BlockAccessor, BlockMetadata, DataBatch, _apply_strict_mode_batch_format
  from ray.types import ObjectRef
  from ray.util.annotations import PublicAPI
label_desc: Use image node_id 7 for calling the DataIterator obj's underlying member method code with example usage: obj.iter_torch_batches() and returns: self, convert_ndarray_batch_to_torch_tensor_batch, batch