code: def run_iter_tensor_batches_benchmark(
      benchmark: Benchmark, data_size_gb: int, block_size_mb: int
  ):
      ctx = ray.data.context.DataContext.get_current()
      ctx.target_max_block_size = block_size_mb * 1024 * 1024
  
      ds = ray.data.read_images(
          f"s3://anonymous@air-example-data-2/{data_size_gb}G-image-data-synthetic-raw"
      )
  
      # Add a label column.
      def add_label(batch):
          label = np.ones(shape=(len(batch), 1))
          batch["label"] = label
          return batch
  
      ds = ds.map_batches(
          add_label, batch_format="pandas"
      ).materialize()
  
      # Test iter_torch_batches() with default args.
      benchmark.run_materialize_ds(
          "iter-torch-batches-default",
          iter_torch_batches,
          ds=ds,
          use_default_params=True,
      )
  
      # Test to_tf() with default args.
      benchmark.run_materialize_ds(
          "to-tf-default",
          to_tf,
          ds=ds,
          feature_columns="image",
          label_columns="label",
          use_default_params=True,
      )
  
      batch_sizes = [16, 32]
  
      # Test with varying batch sizes for iter_torch_batches() and to_tf().
      for batch_size in batch_sizes:
          benchmark.run_materialize_ds(
              f"iter-torch-batches-{batch_size}-block-size-{block_size_mb}",
              iter_torch_batches,
              ds=ds,
              batch_size=batch_size,
          )
  
      prefetch_batches = [0, 1, 4]
      # Test with varying prefetching for iter_torch_batches()
      for prefetch_batch in prefetch_batches:
          for shuffle_buffer_size in [None, 64]:
              test_name = f"iter-torch-batches-bs-{32}-prefetch-{prefetch_batch}-shuffle{shuffle_buffer_size}"  # noqa: E501
              benchmark.run_materialize_ds(
                  test_name,
                  iter_torch_batches,
                  ds=ds,
                  batch_size=32,
                  prefetch_batches=prefetch_batch,
              )
  
      # Test with varying batch sizes and shuffle for iter_torch_batches() and to_tf().
      for batch_size in batch_sizes:
          for shuffle_buffer_size in [batch_size, 2 * batch_size]:
              test_name = f"iter-torch-batches-shuffle-{batch_size}-{shuffle_buffer_size}"
              benchmark.run_materialize_ds(
                  test_name,
                  iter_torch_batches,
                  ds=ds,
                  batch_size=batch_size,
                  local_shuffle_buffer_size=shuffle_buffer_size,
              )
              test_name = (
                  f"to-tf-shuffle-{batch_size}-{shuffle_buffer_size}"
              )
              benchmark.run_materialize_ds(
                  test_name,
                  to_tf,
                  ds=ds,
                  feature_columns="image",
                  label_columns="label",
                  batch_size=batch_size,
                  local_shuffle_buffer_size=shuffle_buffer_size,
              )
  
file_path: ray/release/nightly_tests/dataset/iter_tensor_batches_benchmark.py
imports: import argparse
  import numpy
  from typing import Optional, Union, List
  import ray
  from ray.data.dataset import Dataset
  from benchmark import Benchmark
label_desc: Use image node_id 3 for calling a global function with example usage: run_iter_tensor_batches_benchmark(benchmark, data_size_gb, block_size_mb) and returns: batch