class_docstr: A stage that transforms blocks holistically (e.g., shuffle).
class_docstr_tok: ["A","stage","that","transforms","blocks","holistically","(","e.g.",",","shuffle",")","."]
code: def __init__(
      self,
      name: str,
      num_blocks: Optional[int],
      fn: Callable[[BlockList, bool, Callable], Tuple[BlockList, dict]],
      supports_block_udf: bool = False,
      block_udf: Optional[BlockTransform] = None,
      remote_args: Optional[Dict[str, Any]] = None,
      sub_stage_names: Optional[List[str]] = None,
  ):
      super().__init__(name, num_blocks)
      self.fn = fn
      self.supports_block_udf = supports_block_udf
      self.block_udf = block_udf
      self.ray_remote_args = remote_args or {}
      self.sub_stage_names = sub_stage_names
  
imports: import copy
  import itertools
  from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple, Union
  import ray
  from ray._private.internal_api import get_memory_info_reply, get_state_from_address
  from ray.data._internal.block_list import BlockList
  from ray.data._internal.compute import ActorPoolStrategy, BlockTransform, CallableClass, ComputeStrategy, TaskPoolStrategy, UserDefinedFunction, get_compute, is_task_compute
  from ray.data._internal.dataset_logger import DatasetLogger
  from ray.data._internal.execution.interfaces import TaskContext
  from ray.data._internal.lazy_block_list import LazyBlockList
  from ray.data._internal.logical.interfaces.logical_operator import LogicalOperator
  from ray.data._internal.logical.operators.from_operators import AbstractFrom
  from ray.data._internal.logical.operators.input_data_operator import InputData
  from ray.data._internal.logical.operators.read_operator import Read
  from ray.data._internal.logical.rules.operator_fusion import _are_remote_args_compatible
  from ray.data._internal.stats import DatasetStats, DatasetStatsSummary
  from ray.data._internal.util import create_dataset_tag, unify_block_metadata_schema
  from ray.data.block import Block, BlockMetadata
  from ray.data.context import DataContext
  from ray.types import ObjectRef
  from ray.util.debug import log_once
label_desc: Use image node_id 1 to create a new AllToAllStage object from inherited base classes: Stage with example: obj = AllToAllStage(name, num_blocks, fn, supports_block_udf, block_udf, remote_args, sub_stage_names)