import contextlib import functools import logging import warnings from collections.abc import Callable from typing import Any import torch import torch.utils._pytree as pytree from torch._C import DispatchKey from torch._functorch.utils import exposed_in from torch._higher_order_ops.utils import ( _maybe_compile_and_run_fn, _maybe_run_with_interpreter, check_input_alias_and_mutation_return_outputs, create_bw_fn, create_fn_remove_none, fill_none_with_masks, materialize_as_graph, reenter_make_fx, save_values_for_backward, saved_values, unique_graph_id, validate_subgraph_args_types, ) from torch._ops import HigherOrderOperator from torch._subclasses.fake_tensor import FakeTensorMode from torch.fx.experimental.proxy_tensor import ProxyTorchDispatchMode, track_tensor_tree from torch.utils._python_dispatch import _get_current_dispatch_mode log = logging.getLogger(__name__) class SwitchOp(HigherOrderOperator): def __init__(self): super().__init__("switch") def __call__(self, index, branches, operands): validate_subgraph_args_types(operands) # pyrefly: ignore [missing-attribute] return super().__call__(index, branches, operands) # pyrefly: ignore [bad-override] def gen_schema(self, index, branches, operands): from torch._higher_order_ops.schema import HopSchemaGenerator from torch._higher_order_ops.utils import materialize_as_graph branch_gms: list[torch.fx.GraphModule] = [] branch0_outputs: list[Any] | tuple[Any, ...] = [] mutated_inputs: set[int] = set() for branch in branches: gm = materialize_as_graph(branch, operands) ( _, _, _, branch_mutated_inputs, branch_outputs, ) = check_input_alias_and_mutation_return_outputs(gm) branch_gms.append(gm) if not branch0_outputs: branch0_outputs = branch_outputs mutated_inputs |= set(branch_mutated_inputs) schema_gen = HopSchemaGenerator(self) schema_gen.add_arg("index", index) for i, gm in enumerate(branch_gms): schema_gen.add_arg(f"branch{i}_fn", gm) for idx, arg in enumerate(operands): schema_gen.add_arg(f"operand{idx}", arg, is_mutated=idx in mutated_inputs) for out in branch0_outputs: schema_gen.add_output(out) schema_gen.add_schema_tree_spec(index, branches, operands) return schema_gen.gen_schema() switch_op = SwitchOp() def wrap_branch_fn_flat(*args, branch_fn, spec_operands): operands = pytree.tree_unflatten(args, spec_operands) return branch_fn(*operands) def _get_branch(branches, idx): if not 0 <= idx < len(branches): raise AssertionError( f"switch index {idx} out of range for {len(branches)} branches" ) return branches[idx] @exposed_in("torch") def switch( index: int | torch.SymInt | torch.Tensor, branches: tuple[Callable, ...] | list[Callable], operands: tuple | list = (), ) -> Any: r""" Selects and runs one of N branch functions by index. .. warning:: `torch.switch` is a prototype feature in PyTorch. It has limited support for input and output types. Please look forward to a more stable implementation in a future version of PyTorch. Read more about feature classification at: https://pytorch.org/blog/pytorch-feature-classification-changes/#prototype Equivalent to: ``branches[index](*operands)`` with index in ``[0, len(branches))``. Args: index (Union[int, torch.Tensor]): An int or single-element integer tensor indicating which branch to run. Out-of-range values are clamped into ``[0, len(branches))``. branches (Union[tuple[Callable, ...], list[Callable]]): Non-empty sequence of callables. Each must accept operands and return the same structure of outputs. operands (Tuple of possibly nested dict/list/tuple of torch.Tensor): Inputs to the branch functions. Defaults to (). Restrictions: - Each branch must have the same signature as operands and return the same output structure (shape, dtype, etc.). Constant ``int`` and ``None`` leaves are also permitted in branch outputs and are merged across branches (an unbacked SymInt is introduced when ``int`` leaves differ between branches). - Branches cannot have in-place mutations on inputs or global variables. """ # Flatten operands so the HOP only sees a flat list of tensors. # Each branch is wrapped to unflatten the leaves # back to the user-facing pytree before invocation. leaves_operands, spec_operands = pytree.tree_flatten(operands) def _validate_input(index, branches, operands): if not isinstance(index, (int, torch.SymInt, torch.Tensor)): raise RuntimeError( f"Expected index to be int or single-element tensor, but got {index}." ) if isinstance(index, torch.Tensor) and index.numel() != 1: raise RuntimeError( f"Expected index to be int or single-element tensor, but got {index}." ) if not isinstance(branches, (tuple, list)) or len(branches) == 0: raise RuntimeError( "Expected branches to be a non-empty tuple or list of callables." ) for i, branch in enumerate(branches): if not callable(branch): raise RuntimeError( f"Expected all branches to be callable. branch{i} is not callable." ) for x in operands: if not isinstance(x, torch.Tensor): raise RuntimeError(f"All operand leaves must be a Tensor but got {x}") _validate_input(index, branches, leaves_operands) wrapped_branches = tuple( functools.partial( wrap_branch_fn_flat, branch_fn=branch_fn, spec_operands=spec_operands, ) for branch_fn in branches ) # Early shortcut: single-branch switch degenerates to a plain call num_branches = len(wrapped_branches) if num_branches == 1: return wrapped_branches[0](*leaves_operands) # Clamp out-of-range indices to [0, len(branches) - 1] if isinstance(index, torch.Tensor): index = index.clamp(0, num_branches - 1) elif isinstance(index, torch.SymInt): index = torch.sym_max(0, torch.sym_min(index, num_branches - 1)) elif isinstance(index, int): index = max(0, min(index, num_branches - 1)) # Constant index shortcut for eager mode. if not torch.compiler.is_dynamo_compiling() and isinstance(index, int): # This is the non-strict export case. Strict export and torch.compile are # handled below via _maybe_compile_and_run_fn. if torch.compiler.is_compiling(): warnings.warn( "Index is a Python constant. When used with torch.switch, it specializes on one of the branches." " If you want torch.switch to preserve the branches, please make the index an int tensor or a SymInt.", UserWarning, stacklevel=2, ) return _get_branch(wrapped_branches, index)(*leaves_operands) # Use _maybe_compile_and_run_fn pattern from scan/associative_scan def run_switch(index, wrapped_branches, leaves_operands): return switch_op(index, wrapped_branches, leaves_operands) return _maybe_compile_and_run_fn( run_switch, index, wrapped_branches, leaves_operands ) def trace_switch(proxy_mode, func_overload, index, branches, operands): if not isinstance(operands, (list, tuple)): raise AssertionError( f"Switch operands must be a list or tuple of tensors and SymInts {operands}" ) if not isinstance(branches, (list, tuple)) or len(branches) == 0: raise AssertionError( "Switch branches must be a non-empty list or tuple of callables" ) branch_graphs = [reenter_make_fx(branch)(*operands) for branch in branches] branch_outs = [list() for _ in branches] for i, branch_graph in enumerate(branch_graphs): for node in branch_graph.graph.nodes: if node.op == "output": branch_outs[i].extend(node.args) branch_out_spec = [pytree.tree_flatten(outs)[1] for outs in branch_outs] for i, spec in enumerate(branch_out_spec): if branch_out_spec[0] != spec: raise RuntimeError( "Unmatched output spec from torch.switch branches: " f"branch0 tree_spec {branch_out_spec[0]} vs branch{i} tree_spec {spec}" ) uid, _ = unique_graph_id(proxy_mode, prefix="branch0_graph") for i, branch_graph in enumerate(branch_graphs): proxy_mode.tracer.root.register_module(f"branch{i}_graph_{uid}", branch_graph) args = (index, branch_graphs, operands) proxy_args = pytree.tree_map(proxy_mode.tracer.unwrap_proxy, args) out_proxy = proxy_mode.tracer.create_proxy( "call_function", func_overload, proxy_args, {} ) out = func_overload(index, branch_graphs, operands) return track_tensor_tree(out, out_proxy, constant=None, tracer=proxy_mode.tracer) @switch_op.py_impl(DispatchKey.CompositeExplicitAutograd) def switch_op_dense(index, branches, operands): if not all(isinstance(o, (torch.Tensor, int)) for o in operands): raise RuntimeError( f"Dense implementation operands must be a list of tensors and ints {operands}" ) mode = _get_current_dispatch_mode() if mode is not None: raise AssertionError("Mode should never be enabled for CPU/CUDA key") idx: int = int(index.item()) if isinstance(index, torch.Tensor) else int(index) return _get_branch(branches, idx)(*operands) class SwitchAutogradOp(torch.autograd.Function): @staticmethod # pyrefly: ignore [bad-override] def forward( ctx, index, branches, *operands, ): ctx._index = index # Build one bw fn per branch. ctx._branch_bw_fns = [ create_bw_fn(create_fn_remove_none(branch)[0], operands) for branch in branches ] # We snapshot the dispatch keys in forward for materializing the # bw_graph in backward. ctx._fw_include_key_set = torch._C._dispatch_tls_local_include_set() ctx._fw_exclude_key_set = torch._C._dispatch_tls_local_exclude_set() save_values_for_backward(ctx, operands) with torch._C._AutoDispatchBelowAutograd(): outs = switch_op(index, branches, operands) # Record which output slots are Tensors. Non-Tensor slots (None, # int/SymInt) carry no tangent in backward and the joint is built # to omit them. ctx._fw_output_is_tensor = [ isinstance(o, torch.Tensor) for o in pytree.tree_leaves(outs) ] return outs @staticmethod def backward(ctx, *flat_grads): operands = saved_values(ctx) tensor_grads = tuple( g for g, keep in zip(flat_grads, ctx._fw_output_is_tensor) if keep ) args = operands + tensor_grads # TODO: we need to materialize the bw graphs because dynamo is unable to # trace through the joint function when torch.compile torch.autograd.grad. branches_bw_gm: list[torch.fx.GraphModule] = [] grads_tensor_masks: list[bool] = [] # All branches share the same input signature (see _validate_input) for bw_fn in ctx._branch_bw_fns: wrapped_bw, mask = create_fn_remove_none(bw_fn) bw_gm = materialize_as_graph( wrapped_bw, args, ctx._fw_include_key_set, ctx._fw_exclude_key_set, force_enable_grad=True, ) branches_bw_gm.append(bw_gm) if not grads_tensor_masks: grads_tensor_masks = mask grads = switch_op( ctx._index, branches_bw_gm, args, ) return None, None, *fill_none_with_masks(grads, grads_tensor_masks) @switch_op.py_autograd_impl def switch_autograd(index, branches, operands): return SwitchAutogradOp.apply( index, branches, *operands, ) @switch_op.py_impl(ProxyTorchDispatchMode) def inner(mode, index, branches, operands): return trace_switch(mode, switch_op, index, branches, operands) @switch_op.py_impl(FakeTensorMode) def switch_fake_tensor_mode(mode, index, branches, operands): # Ignore here, because if you've gotten here but you're not manually # tracing the inner graphs, that means that you intend to reuse the graph # directly. Which means the old unbacked symbol bindings are appropriate. # This strategy will not work if unbacked symbols can escape. ignore_fresh_unbacked = contextlib.nullcontext() if mode.shape_env: ignore_fresh_unbacked = mode.shape_env.ignore_fresh_unbacked_symbols() with mode, ignore_fresh_unbacked: flat_branch_outs, branch_out_spec = zip( *[pytree.tree_flatten(branch(*operands)) for branch in branches] ) for i, spec in enumerate(branch_out_spec): if branch_out_spec[0] != spec: raise RuntimeError( "Unmatched output spec from torch.switch branches: " f"branch0 tree_spec {branch_out_spec[0]} vs branch{i} tree_spec {spec}" ) merged_outs = [] for branches_out in zip(*flat_branch_outs): merged_outs.append(_merge_output(branches_out, mode)) return pytree.tree_unflatten(merged_outs, branch_out_spec[0]) def _merge_output(xs: tuple[torch.Tensor | int | None, ...], mode: FakeTensorMode): from torch._higher_order_ops.cond import ( _merge_ints_to_symint, _merge_output as cond_merge_output, ) # Shortcut if a branch produces None outputs; then all branches need to produce None if any(x is None for x in xs): if not all(x is None for x in xs): raise AssertionError(f"expected all leaves to be None, got {xs}") return None # In case all branches return an int, use an unbacked symbol as the merge result if all(type(x) is int for x in xs): return _merge_ints_to_symint(list(xs), mode) # type: ignore[arg-type] return functools.reduce(lambda a, b: cond_merge_output(a, b, mode), xs) @switch_op.py_functionalize_impl def switch_func(ctx, index, branches, inputs): from torch._higher_order_ops.auto_functionalize import ( can_auto_functionalize, do_auto_functionalize_v2, ) from torch._higher_order_ops.utils import _check_alias_and_mutation, HopInstance hop_instance = HopInstance.create(switch_op, index, branches, inputs) if can_auto_functionalize(hop_instance) and hasattr(ctx, "mode"): return do_auto_functionalize_v2( ctx.mode, hop_instance, tuple(pytree.tree_flatten((index, branches, inputs))[0]), {}, ) unwrapped_inputs = ctx.unwrap_tensors(inputs) unwrapped_index = ctx.unwrap_tensors(index) with ctx.redispatch_to_next(): functional_branches = [ ctx.functionalize(_maybe_run_with_interpreter(fn)) for fn in branches ] pre_dispatch = hasattr(ctx, "mode") and ctx.mode.pre_dispatch for i, branch in enumerate(branches): _check_alias_and_mutation( branch, unwrapped_inputs, f"switch_branch{i}", pre_dispatch ) switch_return = switch_op( unwrapped_index, functional_branches, unwrapped_inputs ) return ctx.wrap_tensors(switch_return)