Index

Entropy-guided search tree visualizer.

Public API

• :func:build_entropy_trace -- run the compilation pipeline and return an :class:EntropyTraceBundle wrapping the Rust solver's entropy trace.
• :class:TreeFrameState, :class:TreeStateReducer -- frame-level state derived from a trace for rendering.
• :func:run -- CLI entry point (also accessible via python -m bloqade.lanes.visualize.entropy_tree).

EntropyTraceBundle dataclass

EntropyTraceBundle(
    steps: tuple[TreeTraceStep, ...],
    root_node_id: int,
    best_buffer_size: int,
    arch_spec: Any,
    traced_target: dict[int, LocationAddress],
    blocked_locations: tuple[LocationAddress, ...],
    local_to_global_qid: dict[int, int],
    location_to_global_qid: dict[LocationAddress, int],
    kernel_name: str,
    executed_cz_count: int,
)

Trace plus the metadata the visualizer needs alongside it.

TreeStateReducer

TreeStateReducer(
    steps: tuple[TreeTraceStep, ...] | list[TreeTraceStep],
    root_node_id: int,
    best_buffer_size: int = 2,
)

Converts trace events into renderable tree frames.

Source code in .venv/lib/python3.12/site-packages/bloqade/lanes/visualize/entropy_tree/state.py

def __init__(
    self,
    steps: tuple[TreeTraceStep, ...] | list[TreeTraceStep],
    root_node_id: int,
    best_buffer_size: int = 2,
):
    self.root_node_id = root_node_id
    self.actions = self._expand_actions(list(steps))
    self._display_ids: dict[int, int] = {root_node_id: 0}
    self._next_display_id = 1
    self._root_configuration = self._infer_root_configuration(list(steps))
    self.best_buffer_size = max(0, best_buffer_size)

build_entropy_trace

build_entropy_trace(
    *,
    kernel: Any,
    kernel_name: str,
    layer_index: int = 0,
    max_expansions: int | None = 1000,
    max_goal_candidates: int | None = None
) -> EntropyTraceBundle

Run the compilation pipeline and capture an entropy trace for layer_index.

Drives the full squin->move pipeline with a RustPlacementTraversal configured to collect an entropy trace. Returns all the state the visualizer needs to render the trace.

Source code in .venv/lib/python3.12/site-packages/bloqade/lanes/visualize/entropy_tree/tracer.py

def build_entropy_trace(
    *,
    kernel: Any,
    kernel_name: str,
    layer_index: int = 0,
    max_expansions: int | None = 1000,
    max_goal_candidates: int | None = None,
) -> EntropyTraceBundle:
    """Run the compilation pipeline and capture an entropy trace for ``layer_index``.

    Drives the full squin->move pipeline with a ``RustPlacementTraversal``
    configured to collect an entropy trace. Returns all the state the
    visualizer needs to render the trace.
    """
    from bloqade.analysis import address
    from bloqade.analysis.address.lattice import AddressQubit

    from bloqade.lanes.analysis.layout import LayoutAnalysis
    from bloqade.lanes.arch.gemini.physical import get_arch_spec
    from bloqade.lanes.dialects import move, place
    from bloqade.lanes.heuristics.physical import (
        PhysicalLayoutHeuristicGraphPartitionCenterOut,
        PhysicalPlacementStrategy,
        RustPlacementTraversal,
    )
    from bloqade.lanes.upstream import NativeToPlace, squin_to_move

    arch_spec = get_arch_spec()
    layout_heuristic = PhysicalLayoutHeuristicGraphPartitionCenterOut(
        arch_spec=arch_spec
    )

    goal_candidates = (
        RustPlacementTraversal().max_goal_candidates
        if max_goal_candidates is None
        else max_goal_candidates
    )
    traversal = RustPlacementTraversal(
        strategy="entropy",
        max_expansions=max_expansions,
        max_goal_candidates=goal_candidates,
        collect_entropy_trace=True,
    )
    placement_strategy = PhysicalPlacementStrategy(
        arch_spec=arch_spec,
        traversal=traversal,
    )
    placement_strategy.trace_cz_index = layer_index

    move_main = squin_to_move(
        kernel,
        layout_heuristic=layout_heuristic,
        placement_strategy=placement_strategy,
        no_raise=False,
        logical_initialize=False,
    )
    executed_cz_count = sum(
        1 for stmt in move_main.callable_region.walk() if isinstance(stmt, move.CZ)
    )

    trace = placement_strategy.traced_rust_entropy_trace
    if trace is None:
        raise ValueError(
            f"No entropy trace captured. layer_index={layer_index} may be out of "
            f"range for this kernel ({executed_cz_count} CZ stage(s))."
        )
    steps, root_node_id, best_buffer_size = convert_native_trace(trace)

    traced_target: dict[int, LocationAddress] = placement_strategy.traced_target

    place_main = NativeToPlace(logical_initialize=False).emit(kernel, no_raise=False)
    address_analysis = address.AddressAnalysis(place_main.dialects)
    address_frame, _ = address_analysis.run(place_main)
    all_qubits = tuple(range(address_analysis.next_address))
    initial_layout = LayoutAnalysis(
        place_main.dialects,
        layout_heuristic,
        address_frame.entries,
        all_qubits,
    ).get_layout(place_main)
    location_to_global_qid = {
        location: qid for qid, location in zip(all_qubits, initial_layout, strict=True)
    }
    local_to_global_qid: dict[int, int] = {}
    cz_counter = 0
    for stmt in place_main.callable_region.walk():
        if isinstance(stmt, place.CZ):
            if cz_counter == layer_index:
                owner = stmt.parent_stmt
                while owner is not None and not isinstance(
                    owner, place.StaticPlacement
                ):
                    owner = owner.parent_stmt
                if isinstance(owner, place.StaticPlacement):
                    for local_idx, qubit_ssa in enumerate(owner.qubits):
                        addr = address_frame.entries.get(qubit_ssa)
                        if isinstance(addr, AddressQubit):
                            local_to_global_qid[local_idx] = addr.data
                break
            cz_counter += 1

    return EntropyTraceBundle(
        steps=steps,
        root_node_id=root_node_id,
        best_buffer_size=best_buffer_size,
        arch_spec=arch_spec,
        traced_target=dict(traced_target),
        blocked_locations=placement_strategy.traced_blocked_locations,
        local_to_global_qid=local_to_global_qid,
        location_to_global_qid=location_to_global_qid,
        kernel_name=kernel_name,
        executed_cz_count=executed_cz_count,
    )