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 fuse(self, prev: Stage):
      if not self.can_fuse(prev):
          raise ValueError(
              f"Tried to fuse {prev} with {self}, but these are not fusable."
          )
      assert self.supports_block_udf
      assert (
          prev.fn_constructor_args is None
          and prev.fn_constructor_kwargs is None
      )
      name = prev.name + "->" + self.name
      prev_fn_args = prev.fn_args or tuple()
      prev_fn_args = (
          prev_fn_args if prev.fn is None else (prev.fn,) + prev_fn_args
      )
      prev_fn_kwargs = prev.fn_kwargs or {}
      prev_block_fn = prev.block_fn
      if self.block_udf is None:
  
          def block_udf(
              blocks: Iterable[Block], ctx: TaskContext
          ) -> Iterable[Block]:
              yield from prev_block_fn(
                  blocks, ctx, *prev_fn_args, **prev_fn_kwargs
              )
  
      else:
          self_block_udf = self.block_udf
  
          def block_udf(
              blocks: Iterable[Block], ctx: TaskContext
          ) -> Iterable[Block]:
              blocks = prev_block_fn(
                  blocks,
                  ctx,
                  *prev_fn_args,
                  **prev_fn_kwargs,
              )
              yield from self_block_udf(blocks, ctx)
  
      return AllToAllStage(
          name,
          self.num_blocks,
          self.fn,
          True,
          block_udf,
          prev.ray_remote_args,
          self.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 3 for calling the AllToAllStage obj's underlying member method code with example usage: obj.fuse(prev) and returns: AllToAllStage