Experimental

def __repr__(self) -> str:
    """Return a string representation with compilation statistics."""
    n_direct = len(self._program.direct_f_indices)

    c_graphs = []
    c_params = []
    c_a_terms = []
    c_b_terms = []
    c_c_terms = []
    c_d_terms = []
    c_e_terms = []
    total_memory_bytes = 0
    num_outputs = []
    statevector_widths: list[int] = []
    statevector_actions = 0

    for component in self._program.components:
        if isinstance(component, StatevectorComponent):
            statevector_widths.append(component.width)
            statevector_actions += component.num_actions
            num_outputs.append(len(component.output_indices))
            total_memory_bytes += _nbytes(component)
            continue
        for circuit in component.compiled_scalar_graphs:
            num_outputs.append(len(component.output_indices))
            c_graphs.append(circuit.num_graphs)
            c_params.append(circuit.n_params)
            c_a_terms.append(circuit.node_phases.phases.size)
            c_b_terms.append(circuit.phase_terms.coeffs.size)
            c_c_terms.append(
                circuit.pi_products.psi_params.shape[0]
                * circuit.pi_products.psi_params.shape[1]
            )
            c_d_terms.append(
                circuit.phase_pairs.alpha.size + circuit.phase_pairs.beta.size
            )
            c_e_terms.append(
                circuit.float_terms.table.shape[0]
                * circuit.float_terms.table.shape[1]
            )

            total_memory_bytes += _nbytes(circuit)

    def _format_bytes(n: int) -> str:
        if n < 1024:
            return f"{n} B"
        if n < 1024**2:
            return f"{n / 1024:.1f} kB"
        return f"{n / (1024**2):.1f} MB"

    total_memory_str = _format_bytes(total_memory_bytes)
    error_channel_bits = sum(
        channel.num_bits for channel in self._channel_sampler.channels
    )

    # NOTE: only the statistics of the backends in use are shown.
    parts = [f"{n_direct} direct"]
    if statevector_widths:
        parts.append(
            f"{len(statevector_widths)} statevector components "
            f"(width ≤ {max(statevector_widths)}, "
            f"{statevector_actions} actions)"
        )
    if c_graphs:
        parts.append(f"{sum(c_graphs)} graphs")
    parts += [
        f"{error_channel_bits} error channel bits",
        f"{max(num_outputs, default=0)} outputs for largest cc",
    ]
    if c_graphs:
        parts += [
            f"≤ {max(c_params)} parameters",
            f"{sum(c_a_terms)} A terms",
            f"{sum(c_b_terms)} B terms",
            f"{sum(c_c_terms)} C terms",
            f"{sum(c_d_terms)} D terms",
            f"{sum(c_e_terms)} E terms",
        ]
    parts.append(total_memory_str)
    return f"{type(self).__name__}({', '.join(parts)})"

def probability_of(self, state: np.ndarray, *, batch_size: int) -> np.ndarray:
    """Compute probabilities for a batch of error samples given a measurement state.

    Args:
        state: The measurement outcome state to compute probability for.
        batch_size: Number of error samples to use for estimation.

    Returns:
        Array of probabilities P(state | error_sample) for each error sample.

    """
    if batch_size < 1:
        raise ValueError(f"batch_size must be at least 1, got {batch_size}")
    expected_outputs = self._program.num_outputs
    if state.shape != (expected_outputs,):
        raise ValueError(
            f"state must have shape ({expected_outputs},), got {state.shape}"
        )
    f_samples = jnp.asarray(self._channel_sampler.sample(batch_size))
    p_norm = jnp.ones(batch_size)
    p_joint = jnp.ones(batch_size)

    if len(self._program.direct_f_indices) > 0:
        direct_bits = _direct_bits(self._program, f_samples)
        n_direct = len(self._program.direct_f_indices)
        targets = state[self._program.output_order[:n_direct]]
        p_joint = p_joint * (direct_bits == targets).all(axis=1)

    for component in self._program.components:
        assert isinstance(component, CompiledComponent)
        assert len(component.compiled_scalar_graphs) == 2

        f_selected = f_samples[:, component.f_selection]

        norm_circuit, joint_circuit = component.compiled_scalar_graphs

        # Normalization: only f-params
        p_norm = p_norm * jnp.abs(evaluate(norm_circuit, f_selected))

        # Joint probability: f-params + state
        component_state = state[list(component.output_indices)]
        tiled_state = jnp.tile(component_state, (batch_size, 1))
        joint_params = jnp.hstack([f_selected, tiled_state])
        p_joint = p_joint * jnp.abs(evaluate(joint_circuit, joint_params))

    return np.asarray(p_joint / p_norm)

def diagram(self, **kwargs):
    """Return a timeline-svg diagram of the encoded circuit."""
    return self.circuit.diagram("timeline-svg", **kwargs)

def encode_transversally(self, program_text: str) -> None:
    """Encode a program transversally by replacing  physicalgates with transversal gates.

    Transform a program on m qubits into a program on n * m qubits (consisting of n code blocks).

    Args:
        program_text: The program to encode transversally.

    """
    mod_circ = _transform_circuit(
        program_text,
        stride=self.n,
        offsets=list(range(self.n)),
        stabilizer_generators=self.stabilizer_generators,
        observables=self.observables,
    )
    self.circuit.append_from_stim_program_text(str(mod_circ))

def encoding_flow_generators(self):
    """Return the Pauli flow generators for the encoding circuit."""
    assert self.encoding_program_text is not None
    return stim.Circuit(self.encoding_program_text).flow_generators()

def initialize(
    self, program_text: str, encoding_program_text: str | None = None
) -> None:
    """Initialize state preparation and apply encoding circuit.

    Apply the state preparation program for k qubits, then apply an encoding
    circuit to encode the state into n qubits.

    Args:
        program_text: The state preparation program for k qubits. Generally, this
            should be a simple program that prepares each of the k qubits in a
            single-qubit state.
        encoding_program_text (optional): An encoding circuit for a single logical
            qubit. This should encode a single logical qubit at input
            `self.encoding_qubit` into a state of n qubits.
            If not provided, the encoder will use a noiseless default encoding.

    """
    encoding = encoding_program_text or self.encoding_program_text
    if not encoding:
        raise ValueError("Encoding program text is required")

    mod_circ = _transform_circuit(
        program_text,
        stride=self.n,
        offsets=[self.encoding_qubit],
        used_qubits=self.used_qubits,
        stabilizer_generators=self.stabilizer_generators,
        observables=self.observables,
    )

    self.circuit.append_from_stim_program_text(str(mod_circ))
    self.circuit.append_from_stim_program_text(
        str(
            _transform_circuit(
                encoding,
                stride=1,
                offsets=[self.n * off for off in sorted(self.used_qubits)],
                stabilizer_generators=self.stabilizer_generators,
                observables=self.observables,
            )
        )
    )

def diagram(self, **kwargs):
    """Return a timeline-svg diagram of the encoded circuit."""
    return self.circuit.diagram("timeline-svg", **kwargs)

def encode_transversally(self, program_text: str) -> None:
    """Encode a program transversally by replacing  physicalgates with transversal gates.

    Transform a program on m qubits into a program on n * m qubits (consisting of n code blocks).

    Args:
        program_text: The program to encode transversally.

    """
    mod_circ = _transform_circuit(
        program_text,
        stride=self.n,
        offsets=list(range(self.n)),
        stabilizer_generators=self.stabilizer_generators,
        observables=self.observables,
    )
    self.circuit.append_from_stim_program_text(str(mod_circ))

def encoding_flow_generators(self):
    """Return the Pauli flow generators for the encoding circuit."""
    assert self.encoding_program_text is not None
    return stim.Circuit(self.encoding_program_text).flow_generators()

def initialize(
    self, program_text: str, encoding_program_text: str | None = None
) -> None:
    """Initialize state preparation and apply encoding circuit.

    Apply the state preparation program for k qubits, then apply an encoding
    circuit to encode the state into n qubits.

    Args:
        program_text: The state preparation program for k qubits. Generally, this
            should be a simple program that prepares each of the k qubits in a
            single-qubit state.
        encoding_program_text (optional): An encoding circuit for a single logical
            qubit. This should encode a single logical qubit at input
            `self.encoding_qubit` into a state of n qubits.
            If not provided, the encoder will use a noiseless default encoding.

    """
    encoding = encoding_program_text or self.encoding_program_text
    if not encoding:
        raise ValueError("Encoding program text is required")

    mod_circ = _transform_circuit(
        program_text,
        stride=self.n,
        offsets=[self.encoding_qubit],
        used_qubits=self.used_qubits,
        stabilizer_generators=self.stabilizer_generators,
        observables=self.observables,
    )

    self.circuit.append_from_stim_program_text(str(mod_circ))
    self.circuit.append_from_stim_program_text(
        str(
            _transform_circuit(
                encoding,
                stride=1,
                offsets=[self.n * off for off in sorted(self.used_qubits)],
                stabilizer_generators=self.stabilizer_generators,
                observables=self.observables,
            )
        )
    )