class_docstr: Fuses linear chains of compatible physical operators.
class_docstr_tok: ["Fuses","linear","chains","of","compatible","physical","operators","."]
code: def _get_fused_all_to_all_operator(
      self, down_op: AllToAllOperator, up_op: MapOperator
  ) -> AllToAllOperator:
      assert self._can_fuse(down_op, up_op), (
          "Current rule supports fusing MapOperator -> AllToAllOperator"
          f", but received: {type(up_op).__name__} -> {type(down_op).__name__}"
      )
  
      # Fuse operator names.
      name = up_op.name + "->" + down_op.name
  
      down_logical_op: AbstractAllToAll = self._op_map.pop(down_op)
      up_logical_op: AbstractUDFMap = self._op_map.pop(up_op)
  
      # Fuse transformation functions.
      ray_remote_args = up_logical_op._ray_remote_args
      down_transform_fn = down_op.get_transformation_fn()
      up_map_transformer = up_op.get_map_transformer()
  
      def fused_all_to_all_transform_fn(
          blocks: List[RefBundle], ctx: TaskContext
      ) -> Tuple[List[RefBundle], StatsDict]:
          """To fuse MapOperator->AllToAllOperator, we store the map function
          in the TaskContext so that it may be used by the downstream
          AllToAllOperator's transform function."""
          ctx.upstream_map_transformer = up_map_transformer
          ctx.upstream_map_ray_remote_args = ray_remote_args
          return down_transform_fn(blocks, ctx)
  
      # Make the upstream operator's inputs the new, fused operator's inputs.
      input_deps = up_op.input_dependencies
      assert len(input_deps) == 1
      input_op = input_deps[0]
  
      target_max_block_size = self._get_merged_target_max_block_size(
          up_op.target_max_block_size, down_op.target_max_block_size
      )
  
      op = AllToAllOperator(
          fused_all_to_all_transform_fn,
          input_op,
          target_max_block_size=target_max_block_size,
          num_outputs=down_op._num_outputs,
          # Transfer over the existing sub-progress bars from
          # the AllToAllOperator (if any) into the fused operator.
          sub_progress_bar_names=down_op._sub_progress_bar_names,
          name=name,
      )
      # Bottom out at the source logical op (e.g. Read()).
      input_op = up_logical_op
  
      if isinstance(down_logical_op, RandomShuffle):
          logical_op = RandomShuffle(
              input_op,
              name=name,
              ray_remote_args=ray_remote_args,
          )
      elif isinstance(down_logical_op, Repartition):
          logical_op = Repartition(
              input_op,
              num_outputs=down_logical_op._num_outputs,
              shuffle=down_logical_op._shuffle,
          )
      self._op_map[op] = logical_op
      # Return the fused physical operator.
      return op
  
imports: from typing import List, Optional, Tuple
  from ray.data._internal.compute import get_compute, is_task_compute
  from ray.data._internal.execution.interfaces import PhysicalOperator, RefBundle, TaskContext
  from ray.data._internal.execution.operators.actor_pool_map_operator import ActorPoolMapOperator
  from ray.data._internal.execution.operators.base_physical_operator import AllToAllOperator
  from ray.data._internal.execution.operators.map_operator import MapOperator
  from ray.data._internal.execution.operators.task_pool_map_operator import TaskPoolMapOperator
  from ray.data._internal.logical.interfaces import PhysicalPlan, Rule
  from ray.data._internal.logical.operators.all_to_all_operator import AbstractAllToAll, RandomShuffle, Repartition
  from ray.data._internal.logical.operators.map_operator import AbstractUDFMap
  from ray.data._internal.stats import StatsDict
  from ray.data.context import DataContext
label_desc: Use image node_id 10 for calling the OperatorFusionRule obj's underlying member method code with example usage: obj._get_fused_all_to_all_operator(down_op, up_op) and returns: op, down_transform_fn