src/hamiltonzero/evaluation/runtime.py

# Copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from functools import partial
from pathlib import Path
from typing import Any, NamedTuple

import equinox as eqx
import jax
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

from hamiltonzero.checkpoint import (
    load_model as load_checkpoint_model,
    load_model_metadata,
)
from hamiltonzero.compiled.api import (
    compile_wavefunction,
)
from hamiltonzero.compiled.model import (
    CompiledFinetuneWaveFunction,
    build_finetune_template_model,
)
from hamiltonzero.compiled.types import (
    CompiledWaveFunction,
    CompiledWaveFunctions,
    EnergyInputs,
)
from hamiltonzero.compiled.tree import (
    bind_physical_compiler_kernel,
    compile_physical_tree_from_shared_trunk,
)
from hamiltonzero.compiled.trunk import bind_shared_kernel, compile_shared_trunk
from hamiltonzero.config import EnergyConfig, EvalMCMCConfig, ModelConfig
from hamiltonzero.data.systems import (
    build_context_and_energy,
    load_system as load_spin_hamiltonian,
)
from hamiltonzero.energy import vmc_energy_custom_lap_compiled
from hamiltonzero.energy.custom_lap import build_W_levels
from hamiltonzero.energy.frame import compile_energy_frame, route_energy_inputs
from hamiltonzero.mcmc.runtime import (
    adapt_batched,
    cold_samples,
    init_batched_state,
    run_batched,
)
from hamiltonzero.model.api import build_model
from hamiltonzero.model.context import MultiSystemContext
from hamiltonzero.model.model import _shallow_replace
from hamiltonzero.router.api import (
    route_context as apply_route_context,
    route_state,
)
from hamiltonzero.router.compiled import bind_router_kernel, compile_router_static
from hamiltonzero.router.permutation import (
    permute_ctx_prefix,
    permute_multi_ctx_prefix,
)

from .backend import (
    BeamCandidates,
    CanonicalContext,
    LargeNCompilation,
    MCMCPopulation,
)
from .large_n import compile_large_n_eval_wavefunction


class CompiledEvalContext(eqx.Module):
    mask: Any
    bmask: Any
    route_perm: Any
    energy_frame: Any


class _BeamArrays(NamedTuple):
    permutations: jax.Array
    log_probabilities: jax.Array


def _attention_name(value: str) -> str:
    if value == "tuned":
        return "mhsea_tuned"
    if value == "einsum":
        return "einsum"
    raise ValueError("contextualizer attention must be 'tuned' or 'einsum'")


def _with_contextualizer_attention(model, value: str | None):
    if value is None:
        return model
    contextualizer = model.route_contextualizer
    implementation = _attention_name(value)
    layers = _shallow_replace(
        contextualizer.layers,
        attn_impl=implementation,
    )
    return _shallow_replace(
        model,
        route_contextualizer=_shallow_replace(contextualizer, layers=layers),
    )


def _required_positive_int(metadata: dict[str, Any], name: str) -> int:
    if name not in metadata or isinstance(metadata[name], bool):
        raise ValueError(f"compiled fine-tune checkpoint metadata requires {name!r}")
    value = int(metadata[name])
    if value < 1:
        raise ValueError(
            f"compiled fine-tune checkpoint metadata {name!r} must be positive"
        )
    return value


def _compile_energy_frame(energy_inputs, mask, bmask, permutation):
    return compile_energy_frame(
        energy_inputs,
        mask,
        bmask,
        jnp.asarray(permutation, dtype=jnp.int32),
    )


_compile_frame = jax.jit(_compile_energy_frame)


@jax.jit
def _compile_frame_rows(energy_inputs, mask, bmask, permutations):
    return jax.vmap(
        lambda permutation: _compile_energy_frame(
            energy_inputs,
            mask,
            bmask,
            permutation,
        )
    )(permutations)


def _compiled_wavefunctions_vmap_axes(model: CompiledWaveFunctions):
    return CompiledWaveFunctions(
        kernel=jax.tree_util.tree_map(lambda _value: None, model.kernel),
        trees=jax.tree_util.tree_map(lambda _value: 0, model.trees),
    )


@partial(
    jax.jit,
    static_argnames=("batch_size", "replicas", "initial_m"),
)
def _initialize_single(
    key,
    context,
    initial_sigma,
    *,
    batch_size: int,
    replicas: int,
    initial_m: int,
):
    return init_batched_state(
        jax.random.fold_in(key, jnp.int32(0)),
        context,
        batch_size=batch_size,
        n_replicas=replicas,
        initial_m=initial_m,
        initial_sigma=initial_sigma,
    )


@partial(
    jax.jit,
    static_argnames=("batch_size", "replicas", "initial_m"),
)
def _initialize_rows(
    key,
    context,
    initial_sigma,
    *,
    batch_size: int,
    replicas: int,
    initial_m: int,
):
    indices = jnp.arange(context.mask.shape[0], dtype=jnp.int32)
    return jax.vmap(
        lambda index, context_row: init_batched_state(
            jax.random.fold_in(key, index),
            context_row,
            batch_size=batch_size,
            n_replicas=replicas,
            initial_m=initial_m,
            initial_sigma=initial_sigma,
        )
    )(indices, context)


def _step_single(
    state,
    model,
    context,
    *,
    replica_steps: int,
    walker_chunk_size: int,
):
    return run_batched(
        model,
        context,
        state,
        replica_steps,
        walker_chunk_size=walker_chunk_size,
    )


def _step_compiled_rows(
    state,
    model,
    context,
    *,
    replica_steps: int,
    walker_chunk_size: int,
):
    return jax.vmap(
        lambda state_row, model_row, context_row: run_batched(
            model_row,
            context_row,
            state_row,
            replica_steps,
            walker_chunk_size=walker_chunk_size,
        ),
        in_axes=(0, _compiled_wavefunctions_vmap_axes(model), 0),
    )(state, model, context)


def _adapt_single(
    state,
    beta_history_weight,
    sigma_target,
    sigma_scale,
    haar_target,
):
    return adapt_batched(
        state,
        beta_history_weight=beta_history_weight,
        sigma_target=sigma_target,
        sigma_scale=sigma_scale,
        haar_target=haar_target,
    )


def _adapt_rows(
    state,
    beta_history_weight,
    sigma_target,
    sigma_scale,
    haar_target,
):
    return jax.vmap(
        lambda row: adapt_batched(
            row,
            beta_history_weight=beta_history_weight,
            sigma_target=sigma_target,
            sigma_scale=sigma_scale,
            haar_target=haar_target,
        )
    )(state)


def _compiled_energy_single(kernel, tree, frame, q, *, chunk_size: int):
    return vmc_energy_custom_lap_compiled(
        kernel,
        tree,
        frame,
        q,
        chunk_size=chunk_size,
    )


def _build_compiled_energy_single(mesh: Mesh, chunk_size: int):
    batch = P("batch", None, None)
    batch_output = P("batch")

    def local_energy(kernel, tree, frame, q):
        outputs = _compiled_energy_single(
            kernel,
            tree,
            frame,
            q,
            chunk_size=int(chunk_size),
        )
        local_count = jnp.asarray(q.shape[0], dtype=jnp.int32)
        global_count = jax.lax.psum(local_count, "batch")
        guard = global_count.astype(jnp.float32) * jnp.asarray(0.0, jnp.float32)
        return tuple(value + guard.astype(value.dtype) for value in outputs)

    mapped = jax.shard_map(
        local_energy,
        mesh=mesh,
        in_specs=(P(), P(), P(), batch),
        out_specs=(batch_output,) * 4,
        check_vma=False,
    )
    replicated_sharding = NamedSharding(mesh, P())
    batch_sharding = NamedSharding(mesh, batch)
    batch_output_sharding = NamedSharding(mesh, batch_output)
    return jax.jit(
        mapped,
        in_shardings=(
            replicated_sharding,
            replicated_sharding,
            replicated_sharding,
            batch_sharding,
        ),
        out_shardings=(batch_output_sharding,) * 4,
    )


def _compiled_energy_rows(kernel, trees, frames, q, *, chunk_size: int):
    def energy_row(tree, frame, q_row):
        n_sites = int(q_row.shape[-2])
        frame = eqx.tree_at(
            lambda value: value.w_levels,
            frame,
            tuple(build_W_levels(frame.custom_lap_J_eff, n_sites)),
        )
        return vmc_energy_custom_lap_compiled(
            kernel,
            tree,
            frame,
            q_row,
            chunk_size=chunk_size,
        )

    return jax.vmap(energy_row)(trees, frames, q)


def _build_compiled_energy_rows(mesh: Mesh, chunk_size: int):
    systems = P("systems")
    system_batch = P("systems", None)

    def local_energy(kernel, trees, frames, q):
        outputs = _compiled_energy_rows(
            kernel,
            trees,
            frames,
            q,
            chunk_size=int(chunk_size),
        )
        local_count = jnp.asarray(q.shape[0] * q.shape[1], dtype=jnp.int32)
        global_count = jax.lax.psum(local_count, "systems")
        guard = global_count.astype(jnp.float32) * jnp.asarray(0.0, jnp.float32)
        return tuple(value + guard.astype(value.dtype) for value in outputs)

    mapped = jax.shard_map(
        local_energy,
        mesh=mesh,
        in_specs=(P(), systems, systems, system_batch),
        out_specs=(system_batch,) * 4,
        check_vma=False,
    )
    replicated_sharding = NamedSharding(mesh, P())
    systems_sharding = NamedSharding(mesh, systems)
    system_batch_sharding = NamedSharding(mesh, system_batch)
    return jax.jit(
        mapped,
        in_shardings=(
            replicated_sharding,
            systems_sharding,
            systems_sharding,
            system_batch_sharding,
        ),
        out_shardings=(system_batch_sharding,) * 4,
    )


_compile_single = jax.jit(compile_wavefunction)


def _systems_mesh(count: int) -> Mesh:
    devices = tuple(jax.devices())
    lanes = min(len(devices), int(count))
    if int(count) % lanes:
        raise ValueError(
            f"contest K={count} must be divisible by visible devices={lanes}"
        )
    return Mesh(np.asarray(devices[:lanes], dtype=object), ("systems",))


def _batch_mesh(batch_size: int) -> Mesh:
    devices = tuple(jax.devices())
    lanes = min(len(devices), int(batch_size))
    while int(batch_size) % lanes:
        lanes -= 1
    return Mesh(np.asarray(devices[:lanes], dtype=object), ("batch",))


def _single_state_sharding(mesh: Mesh, state):
    walkers = NamedSharding(mesh, P("batch"))
    replicated = NamedSharding(mesh, P())
    return type(state)(
        q=walkers,
        log_p=walkers,
        grad_log_p=walkers,
        beta=replicated,
        sigma=replicated,
        step=replicated,
        key=walkers,
        n_local_accept=walkers,
        n_local=walkers,
        n_swap_accept=walkers,
        n_swap=walkers,
        mask=replicated,
        m=replicated,
        n_haar_accept=walkers,
        n_haar=walkers,
    )


def _single_state_specs(state):
    walkers = P("batch")
    replicated = P()
    return type(state)(
        q=walkers,
        log_p=walkers,
        grad_log_p=walkers,
        beta=replicated,
        sigma=replicated,
        step=replicated,
        key=walkers,
        n_local_accept=walkers,
        n_local=walkers,
        n_swap_accept=walkers,
        n_swap=walkers,
        mask=replicated,
        m=replicated,
        n_haar_accept=walkers,
        n_haar=walkers,
    )


def _row_state_specs(state):
    systems = P("systems")
    return jax.tree_util.tree_map(lambda _value: systems, state)


def _row_model_specs(model):
    return CompiledWaveFunctions(
        kernel=jax.tree_util.tree_map(lambda _value: P(), model.kernel),
        trees=jax.tree_util.tree_map(lambda _value: P("systems"), model.trees),
    )


def _row_model_sharding(mesh: Mesh, model):
    return CompiledWaveFunctions(
        kernel=_replicated_sharding(mesh, model.kernel),
        trees=_system_sharding(mesh, model.trees),
    )


def _build_rows_mcmc_step(
    mesh: Mesh,
    state,
    model,
    context,
    *,
    replica_steps: int,
    walker_chunk_size: int,
):
    state_specs = _row_state_specs(state)
    model_specs = _row_model_specs(model)
    context_specs = jax.tree_util.tree_map(lambda _value: P("systems"), context)
    state_sharding = _system_sharding(mesh, state)
    model_sharding = _row_model_sharding(mesh, model)
    context_sharding = _system_sharding(mesh, context)

    def local_step(state_local, model_local, context_local):
        out = _step_compiled_rows(
            state_local,
            model_local,
            context_local,
            replica_steps=int(replica_steps),
            walker_chunk_size=int(walker_chunk_size),
        )
        local_count = jnp.asarray(out.q.shape[0] * out.q.shape[1], dtype=jnp.int32)
        global_count = jax.lax.psum(local_count, "systems")
        guard = global_count.astype(out.q.dtype) * jnp.asarray(0.0, out.q.dtype)
        return eqx.tree_at(lambda value: value.q, out, out.q + guard)

    mapped = jax.shard_map(
        local_step,
        mesh=mesh,
        in_specs=(state_specs, model_specs, context_specs),
        out_specs=state_specs,
        check_vma=False,
    )
    return jax.jit(
        mapped,
        in_shardings=(state_sharding, model_sharding, context_sharding),
        out_shardings=state_sharding,
        donate_argnums=(0,),
    )


def _build_singular_mcmc_step(
    mesh: Mesh,
    state,
    state_sharding,
    model_sharding,
    context_sharding,
    *,
    replica_steps: int,
    walker_chunk_size: int,
):
    state_specs = _single_state_specs(state)

    def local_step(state_local, model_local, context_local):
        out = _step_single(
            state_local,
            model_local,
            context_local,
            replica_steps=int(replica_steps),
            walker_chunk_size=int(walker_chunk_size),
        )
        local_count = jnp.asarray(out.q.shape[0], dtype=jnp.int32)
        global_count = jax.lax.psum(local_count, "batch")
        guard = global_count.astype(out.q.dtype) * jnp.asarray(0.0, out.q.dtype)
        return eqx.tree_at(lambda value: value.q, out, out.q + guard)

    mapped = jax.shard_map(
        local_step,
        mesh=mesh,
        in_specs=(state_specs, P(), P()),
        out_specs=state_specs,
        check_vma=False,
    )
    return jax.jit(
        mapped,
        in_shardings=(state_sharding, model_sharding, context_sharding),
        out_shardings=state_sharding,
        donate_argnums=(0,),
    )


def _system_sharding(mesh: Mesh, value):
    return jax.tree_util.tree_map(
        lambda array: NamedSharding(
            mesh,
            P("systems", *([None] * (array.ndim - 1))),
        ),
        value,
    )


def _replicated_sharding(mesh: Mesh, value):
    replicated = NamedSharding(mesh, P())
    return jax.tree_util.tree_map(lambda _array: replicated, value)


def _place_system_rows(mesh: Mesh, value):
    return jax.device_put(value, _system_sharding(mesh, value))


def _abstract(tree):
    return jax.tree_util.tree_map(
        lambda value: jax.ShapeDtypeStruct(value.shape, value.dtype),
        tree,
    )


def _beam_mesh(width: int) -> Mesh:
    devices = tuple(jax.devices())
    lanes = len(devices) if int(width) % len(devices) == 0 else 1
    return Mesh(np.asarray(devices[:lanes], dtype=object), ("systems",))


def _distributed_beam_local(decoder, static, tau, *, width: int, lanes: int):
    permutations, log_probabilities = decoder.beam_search(
        static.node_input,
        static.raw_edge,
        static.routable_mask,
        global_feat=static.global_input,
        tau=tau,
        beam_width=width,
        real_mask=static.real_mask,
        first_orbit_ids=(
            static.quotient_node_key,
            static.quotient_edge_key,
            static.needs_fwl2,
        ),
        router_static=static,
        distributed_axis_name="systems",
        distributed_lanes=lanes,
    )
    return _BeamArrays(permutations, log_probabilities)


def _build_distributed_beam(mesh: Mesh, decoder, static, width: int):
    lanes = int(mesh.shape["systems"])
    if tuple(mesh.axis_names) != ("systems",) or int(width) % lanes:
        raise ValueError("distributed eval beam requires a divisible systems mesh")
    mapped = jax.shard_map(
        partial(
            _distributed_beam_local,
            width=int(width),
            lanes=lanes,
        ),
        mesh=mesh,
        in_specs=(
            jax.tree_util.tree_map(lambda _leaf: P(), decoder),
            jax.tree_util.tree_map(lambda _leaf: P(), static),
            P(),
        ),
        out_specs=_BeamArrays(P(), P()),
        check_vma=False,
    )
    replicated = NamedSharding(mesh, P())
    return jax.jit(
        mapped,
        in_shardings=(
            jax.tree_util.tree_map(lambda _leaf: replicated, decoder),
            jax.tree_util.tree_map(lambda _leaf: replicated, static),
            replicated,
        ),
        out_shardings=_BeamArrays(
            replicated,
            replicated,
        ),
    )


def _compile_eval_router_static(model, context):
    trunk = compile_shared_trunk(model, context)
    kernel = bind_router_kernel(model)
    return compile_router_static(
        kernel,
        trunk,
        context.route_quotient_node_key,
        context.route_quotient_edge_key,
        context.needs_fwl2,
    )


class DefaultEvalBackend:
    def __init__(self) -> None:
        self._energy_inputs_by_context: dict[int, EnergyInputs] = {}
        self._energy_frames: dict[int, Any] = {}
        self._context_meshes: dict[int, Mesh] = {}
        self._contest_mesh: Mesh | None = None
        self._contest_step_entries: dict[tuple[int, int], Any] = {}
        self._contest_adapt_entry: Any | None = None
        self._singular_mesh: Mesh | None = None
        self._singular_state_sharding: Any | None = None
        self._singular_model_sharding: Any | None = None
        self._singular_context_sharding: Any | None = None
        self._singular_step_entries: dict[tuple[int, int], Any] = {}
        self._singular_adapt_entry: Any | None = None
        self._singular_energy_entries: dict[int, Any] = {}
        self._contest_energy_entries: dict[int, Any] = {}

    def build_system(self, system, energy: EnergyConfig):
        context, energy_inputs = build_context_and_energy(
            system,
            n_max=None,
            mu=energy.mu,
            eps=energy.eps,
        )
        self._energy_inputs_by_context[id(context)] = energy_inputs
        return context

    def load_system(self, path: Path, energy: EnergyConfig):
        return self.build_system(load_spin_hamiltonian(path), energy)

    def load_model(
        self,
        checkpoint: Path,
        config: ModelConfig,
        key,
        context,
        *,
        contextualizer_attention: str | None,
    ):
        metadata = load_model_metadata(checkpoint) or {}
        kind = metadata.get("kind", "router")
        n_sites = int(context.mask.shape[-1])
        eager_template = build_model(config, key, n_max=n_sites)
        if kind == "router":
            eager_template = _with_contextualizer_attention(
                eager_template,
                contextualizer_attention,
            )
            return load_checkpoint_model(checkpoint, eager_template)
        if kind != "compiled_finetune":
            raise ValueError(f"unsupported checkpoint kind {kind!r}")
        leaf_rank = _required_positive_int(metadata, "leaf_rank")
        merge_rank = _required_positive_int(metadata, "merge_rank")
        checkpoint_n = int(metadata.get("n_max", n_sites))
        if checkpoint_n != n_sites:
            raise ValueError(
                "compiled fine-tune checkpoint width does not match the "
                f"evaluation system: checkpoint={checkpoint_n}, system={n_sites}"
            )
        template = build_finetune_template_model(
            eager_template,
            n_sites,
            leaf_rank=leaf_rank,
            merge_rank=merge_rank,
        )
        return load_checkpoint_model(checkpoint, template)

    def canonicalize_context(self, context) -> CanonicalContext:
        route = jnp.asarray(context.route_perm, dtype=jnp.int32)
        if route.ndim != 1:
            raise ValueError("single-system context route must have shape [N]")
        inverse = jnp.argsort(route).astype(jnp.int32)
        canonical = permute_ctx_prefix(context, inverse)
        identity = jnp.arange(route.shape[0], dtype=jnp.int32)
        canonical = eqx.tree_at(
            lambda value: value.route_perm,
            canonical,
            identity,
        )
        energy_inputs = self._energy_inputs_by_context.get(id(context))
        if energy_inputs is None:
            raise RuntimeError("energy inputs are unavailable for this context")
        self._energy_inputs_by_context[id(canonical)] = route_energy_inputs(
            energy_inputs,
            inverse,
        )
        return CanonicalContext(
            context=canonical,
            old_inverse=inverse[None, :],
        )

    def embedded_route(self, model):
        if isinstance(model, CompiledFinetuneWaveFunction):
            return jnp.asarray(model.perm, dtype=jnp.int32)
        return None

    def route_context(
        self,
        context,
        permutation,
        *,
        compact_custom_lap: bool,
    ):
        permutation = jnp.asarray(permutation, dtype=jnp.int32)
        if permutation.ndim == 2:
            if permutation.shape[0] != 1:
                raise ValueError("single-system route must have shape [1, N]")
            permutation = permutation[0]
        if permutation.ndim != 1:
            raise ValueError("single-system route must have shape [N]")
        energy_inputs = self._energy_inputs_by_context.get(id(context))
        if energy_inputs is None:
            raise RuntimeError("energy inputs are unavailable for this context")
        frame = _compile_frame(
            energy_inputs,
            context.mask,
            context.bmask,
            permutation,
        )
        if compact_custom_lap:
            routed = CompiledEvalContext(
                mask=frame.masks.real,
                bmask=frame.masks.balanced,
                route_perm=permutation,
                energy_frame=frame,
            )
            self._energy_frames[id(routed)] = frame
            return routed
        routed = apply_route_context(context, permutation)
        self._energy_frames[id(routed)] = frame
        return routed

    def virtual_context(self, context, permutations):
        permutations = jnp.asarray(permutations, dtype=jnp.int32)
        if permutations.ndim != 2:
            raise ValueError("candidate permutations must have shape [K, N]")
        count = permutations.shape[0]
        batched = MultiSystemContext.from_single(context)
        tiled = jax.tree_util.tree_map(
            lambda value: (
                jnp.repeat(value, count, axis=0)
                if eqx.is_array(value) and value.ndim >= 1
                else value
            ),
            batched,
        )
        tiled = eqx.tree_at(
            lambda value: value.route_perm,
            tiled,
            permutations,
        )
        routed = permute_multi_ctx_prefix(tiled, permutations)
        energy_inputs = self._energy_inputs_by_context.get(id(context))
        if energy_inputs is None:
            raise RuntimeError("energy inputs are unavailable for this context")
        frames = _compile_frame_rows(
            energy_inputs,
            context.mask,
            context.bmask,
            permutations,
        )
        mesh = _systems_mesh(count)
        routed = _place_system_rows(mesh, routed)
        frames = _place_system_rows(mesh, frames)
        self._energy_frames[id(routed)] = frames
        self._context_meshes[id(routed)] = mesh
        self._contest_mesh = mesh
        self._contest_step_entries.clear()
        self._contest_adapt_entry = None
        self._contest_energy_entries.clear()
        return routed

    def release_context(self, context) -> None:
        self._energy_inputs_by_context.pop(id(context), None)
        self._energy_frames.pop(id(context), None)
        self._context_meshes.pop(id(context), None)
        self._contest_mesh = None
        self._contest_step_entries.clear()
        self._contest_adapt_entry = None
        self._contest_energy_entries.clear()

    def beam_candidates(
        self,
        model,
        context,
        *,
        beam_width: int,
        top_k: int,
        temperature: float,
    ) -> BeamCandidates:
        decoder = getattr(model, "route_decoder", None)
        from hamiltonzero.model.route_pointer import TreePrefixPointerMHSEA

        if not isinstance(decoder, TreePrefixPointerMHSEA):
            raise ValueError("eval requires the learned-quotient TreePrefix decoder")
        if int(top_k) > int(beam_width):
            raise ValueError("top_k cannot exceed beam_width")
        mesh = _beam_mesh(int(beam_width))
        static = eqx.filter_jit(_compile_eval_router_static)(model, context)
        replicated = NamedSharding(mesh, P())
        decoder, static, tau = jax.device_put(
            (
                decoder,
                static,
                jnp.asarray(temperature, dtype=jnp.float32),
            ),
            replicated,
        )
        result = _build_distributed_beam(
            mesh,
            decoder,
            static,
            int(beam_width),
        )(decoder, static, tau)
        jax.block_until_ready(result.permutations)
        permutations = result.permutations[: int(top_k)].astype(jnp.int32)
        log_probabilities = result.log_probabilities[: int(top_k)].astype(jnp.float32)
        return BeamCandidates(
            permutations=permutations[None],
            log_probabilities=log_probabilities[None],
        )

    def compile_single(self, model, routed_context):
        n_sites = int(routed_context.mask.shape[-1])
        return _compile_single(
            model,
            routed_context,
            jnp.arange(n_sites, dtype=jnp.int32),
        )

    def compile_embedded(self, model):
        if not isinstance(model, CompiledFinetuneWaveFunction):
            raise TypeError("embedded eval compilation requires a fine-tune checkpoint")
        return CompiledWaveFunction(
            kernel=model.kernel,
            tree=model.as_compiled_tree(),
        )

    def compile_candidates(self, model, canonical_context, permutations):
        if self._contest_mesh is None:
            raise RuntimeError("contest context must be built before compilation")
        mesh = self._contest_mesh
        physical_kernel = bind_physical_compiler_kernel(model)
        shared_trunk = jax.jit(compile_shared_trunk)(model, canonical_context)
        jax.block_until_ready(shared_trunk)
        physical_sharding = _replicated_sharding(mesh, physical_kernel)
        trunk_sharding = _replicated_sharding(mesh, shared_trunk)
        permutation_sharding = NamedSharding(mesh, P("systems", None))
        physical_kernel = jax.device_put(physical_kernel, physical_sharding)
        shared_trunk = jax.device_put(shared_trunk, trunk_sharding)
        permutations = jax.device_put(
            jnp.asarray(permutations, dtype=jnp.int32),
            permutation_sharding,
        )

        def compile_all(kernel, trunk, candidate_permutations):
            return jax.vmap(
                lambda permutation: compile_physical_tree_from_shared_trunk(
                    kernel,
                    trunk,
                    permutation,
                )
            )(candidate_permutations)

        tree_template = jax.eval_shape(
            compile_all,
            _abstract(physical_kernel),
            _abstract(shared_trunk),
            _abstract(permutations),
        )
        tree_sharding = _system_sharding(mesh, tree_template)
        trees = jax.jit(
            compile_all,
            in_shardings=(
                physical_sharding,
                trunk_sharding,
                permutation_sharding,
            ),
            out_shardings=tree_sharding,
        )(physical_kernel, shared_trunk, permutations)
        shared_kernel = bind_shared_kernel(model)
        shared_kernel = jax.device_put(
            shared_kernel,
            _replicated_sharding(mesh, shared_kernel),
        )
        compiled = CompiledWaveFunctions(shared_kernel, trees)
        jax.block_until_ready(compiled)
        return compiled

    def select_candidate(self, wavefunctions, winner: int):
        if self._contest_mesh is None:
            raise RuntimeError("contest mesh is unavailable for winner selection")
        mesh = self._contest_mesh
        count = int(wavefunctions.trees.perm.shape[0])
        index = int(winner)
        if index < 0 or index >= count:
            raise IndexError(
                f"winner index {index} outside candidate range [0, {count})"
            )
        tree_sharding = _system_sharding(mesh, wavefunctions.trees)
        winner_sharding = NamedSharding(mesh, P())

        def gather(trees, selected_index):
            return jax.tree_util.tree_map(
                lambda value: jax.lax.dynamic_index_in_dim(
                    value,
                    selected_index,
                    axis=0,
                    keepdims=False,
                ),
                trees,
            )

        output_template = jax.eval_shape(
            gather,
            _abstract(wavefunctions.trees),
            jax.ShapeDtypeStruct((), jnp.int32),
        )
        tree = jax.jit(
            gather,
            in_shardings=(tree_sharding, winner_sharding),
            out_shardings=_replicated_sharding(mesh, output_template),
        )(
            wavefunctions.trees,
            jax.device_put(jnp.asarray(index, jnp.int32), winner_sharding),
        )
        selected = CompiledWaveFunction(wavefunctions.kernel, tree)
        jax.block_until_ready(selected)
        return selected

    def compile_large_n(
        self,
        model,
        canonical_context,
        *,
        sequence_shards: int,
        pair_tile_size: int,
        temperature: float,
    ) -> LargeNCompilation:
        result = compile_large_n_eval_wavefunction(
            model,
            canonical_context,
            seq_shards=int(sequence_shards),
            pair_tile_size=int(pair_tile_size),
            tau=float(temperature),
        )
        device = jax.devices()[0]
        return LargeNCompilation(
            wavefunction=jax.device_put(result.wavefunction, device),
            permutation=jax.device_put(result.perm, device),
            log_probability=jax.device_put(result.logp, device),
        )

    def prepare_singular(self, model, context, state):
        if state.q.ndim != 4:
            raise ValueError(
                "post-selection eval MCMC state must have shape [B, R, N, 4]"
            )
        frame = self._frames(context)
        mesh = _batch_mesh(int(state.q.shape[0]))
        state_sharding = _single_state_sharding(mesh, state)
        model_sharding = _replicated_sharding(mesh, model)
        context_sharding = _replicated_sharding(mesh, context)
        model = jax.device_put(model, model_sharding)
        context = jax.device_put(context, context_sharding)
        state = jax.device_put(state, state_sharding)
        self._energy_frames[id(context)] = frame
        jax.block_until_ready((model, context, state.q))
        self._singular_mesh = mesh
        self._singular_state_sharding = state_sharding
        self._singular_model_sharding = model_sharding
        self._singular_context_sharding = context_sharding
        self._singular_step_entries.clear()
        self._singular_adapt_entry = None
        self._singular_energy_entries.clear()
        return model, context, state

    def initialize_mcmc(
        self,
        key,
        model,
        context,
        config: EvalMCMCConfig,
    ):
        del model
        if isinstance(context, MultiSystemContext):
            state = _initialize_rows(
                key,
                context,
                jnp.asarray(config.initial_sigma, dtype=jnp.float32),
                batch_size=int(config.batch_size),
                replicas=int(config.replicas),
                initial_m=int(config.initial_haar_sites),
            )
            mesh = self._context_meshes.get(id(context))
            return _place_system_rows(mesh, state) if mesh is not None else state
        return _initialize_single(
            key,
            context,
            batch_size=int(config.batch_size),
            replicas=int(config.replicas),
            initial_m=int(config.initial_haar_sites),
            initial_sigma=jnp.asarray(config.initial_sigma, dtype=jnp.float32),
        )

    def step_mcmc(
        self,
        state,
        model,
        context,
        *,
        replica_steps: int,
        walker_chunk_size: int,
    ):
        if state.q.ndim != 5:
            if (
                self._singular_mesh is None
                or self._singular_state_sharding is None
                or self._singular_model_sharding is None
                or self._singular_context_sharding is None
            ):
                raise RuntimeError("singular eval placement has not been prepared")
            key = (int(replica_steps), int(walker_chunk_size))
            step = self._singular_step_entries.get(key)
            if step is None:
                step = _build_singular_mcmc_step(
                    self._singular_mesh,
                    state,
                    self._singular_state_sharding,
                    self._singular_model_sharding,
                    self._singular_context_sharding,
                    replica_steps=key[0],
                    walker_chunk_size=key[1],
                )
                self._singular_step_entries[key] = step
            return step(state, model, context)
        if not isinstance(model, CompiledWaveFunctions):
            raise TypeError("multirow eval MCMC requires compiled wavefunctions")
        mesh = self._context_meshes.get(id(context))
        if mesh is None:
            raise RuntimeError("contest MCMC mesh is unavailable")
        key = (int(replica_steps), int(walker_chunk_size))
        step = self._contest_step_entries.get(key)
        if step is None:
            step = _build_rows_mcmc_step(
                mesh,
                state,
                model,
                context,
                replica_steps=key[0],
                walker_chunk_size=key[1],
            )
            self._contest_step_entries[key] = step
        return step(state, model, context)

    def adapt_mcmc(self, state, config: EvalMCMCConfig):
        arguments = (
            state,
            jnp.asarray(config.beta_history_weight, dtype=jnp.float32),
            jnp.asarray(config.langevin_target_acceptance, dtype=jnp.float32),
            jnp.asarray(config.sigma_scale, dtype=jnp.float32),
            jnp.asarray(config.haar_target_acceptance, dtype=jnp.float32),
        )
        if state.q.ndim == 5:
            if self._contest_mesh is None:
                raise RuntimeError("contest adaptation mesh is unavailable")
            if self._contest_adapt_entry is None:
                replicated = NamedSharding(self._contest_mesh, P())
                state_sharding = _system_sharding(self._contest_mesh, state)
                self._contest_adapt_entry = jax.jit(
                    _adapt_rows,
                    in_shardings=(
                        state_sharding,
                        replicated,
                        replicated,
                        replicated,
                        replicated,
                    ),
                    out_shardings=state_sharding,
                )
            return self._contest_adapt_entry(*arguments)
        if self._singular_mesh is None or self._singular_state_sharding is None:
            raise RuntimeError("singular eval placement has not been prepared")
        if self._singular_adapt_entry is None:
            replicated = NamedSharding(self._singular_mesh, P())
            self._singular_adapt_entry = jax.jit(
                _adapt_single,
                in_shardings=(
                    self._singular_state_sharding,
                    replicated,
                    replicated,
                    replicated,
                    replicated,
                ),
                out_shardings=self._singular_state_sharding,
            )
        return self._singular_adapt_entry(*arguments)

    def route_mcmc(self, state, permutation):
        return route_state(state, permutation)

    def mcmc_population(self, state) -> MCMCPopulation:
        return MCMCPopulation(q=state.q, sigma=state.sigma, beta=state.beta)

    def replace_mcmc_population(
        self,
        state,
        population: MCMCPopulation,
    ):
        q = population.q
        sigma = population.sigma
        beta = population.beta
        if q.ndim == state.q.ndim + 1 and q.shape[0] == 1:
            q = jax.device_put(q[0], jax.devices()[0])
        if sigma.ndim == state.sigma.ndim + 1 and sigma.shape[0] == 1:
            sigma = jax.device_put(sigma[0], jax.devices()[0])
        if beta.ndim == state.beta.ndim + 1 and beta.shape[0] == 1:
            beta = jax.device_put(beta[0], jax.devices()[0])
        return eqx.tree_at(
            lambda value: (value.q, value.sigma, value.beta),
            state,
            (
                q.astype(state.q.dtype),
                sigma.astype(state.sigma.dtype),
                beta.astype(state.beta.dtype),
            ),
        )

    def cold_walkers(self, state):
        return (
            jax.vmap(cold_samples)(state) if state.q.ndim == 5 else cold_samples(state)
        )

    def _frames(self, context):
        if isinstance(context, CompiledEvalContext):
            return context.energy_frame
        frames = self._energy_frames.get(id(context))
        if frames is not None:
            return frames
        energy_inputs = self._energy_inputs_by_context.get(id(context))
        if energy_inputs is None:
            raise RuntimeError("energy frame is unavailable for this context")
        if isinstance(context, MultiSystemContext):
            raise RuntimeError("multi-system energy frames must be compiled explicitly")
        n_sites = int(context.mask.shape[-1])
        frames = _compile_frame(
            energy_inputs,
            context.mask,
            context.bmask,
            jnp.arange(n_sites, dtype=jnp.int32),
        )
        self._energy_frames[id(context)] = frames
        return frames

    def custom_lap_energy(
        self,
        model,
        context,
        q,
        config: EnergyConfig,
    ):
        chunk_size = int(config.chunk_size)
        if isinstance(model, CompiledWaveFunctions):
            frames = self._frames(context)
            mesh = self._context_meshes.get(id(context))
            if mesh is None:
                raise RuntimeError("contest energy mesh is unavailable")
            energy = self._contest_energy_entries.get(chunk_size)
            if energy is None:
                energy = _build_compiled_energy_rows(mesh, chunk_size)
                self._contest_energy_entries[chunk_size] = energy
            return energy(
                model.kernel,
                model.trees,
                frames,
                q,
            )
        if not isinstance(model, CompiledWaveFunction):
            raise TypeError("singular eval energy requires a compiled wavefunction")
        if (
            self._singular_mesh is None
            or self._singular_model_sharding is None
            or self._singular_context_sharding is None
        ):
            raise RuntimeError("singular eval placement has not been prepared")
        q_sharding = NamedSharding(
            self._singular_mesh,
            P("batch", None, None),
        )
        q = jax.device_put(q, q_sharding)
        frames = self._frames(context)
        frame_sharding = _replicated_sharding(
            self._singular_mesh,
            frames,
        )
        frames = jax.device_put(frames, frame_sharding)
        energy = self._singular_energy_entries.get(chunk_size)
        if energy is None:
            energy = _build_compiled_energy_single(
                self._singular_mesh,
                chunk_size,
            )
            self._singular_energy_entries[chunk_size] = energy
        outputs = energy(
            model.kernel,
            model.tree,
            frames,
            q,
        )
        return tuple(jnp.expand_dims(value, axis=0) for value in outputs)

    def block_until_ready(self, value) -> None:
        jax.block_until_ready(value)


def build_eval_backend() -> DefaultEvalBackend:
    return DefaultEvalBackend()


__all__ = [
    "DefaultEvalBackend",
    "build_eval_backend",
]