quchip.chip.dressing

JAX-traceable dressed-state labeling primitives.

Three layers, separated cleanly:

1. Pure-JAX kernel. label_eigensystem() takes (evals, evecs, reference, policy) and returns a Labeling. Knows nothing about Chip, devices, caching, warnings, or backends.

2. References as data. BareProductReference and EigenstateReference produce reference vectors and label keys. Overlaps are computed by one shared kernel, compute_overlaps(), so subspace generalization is a tensor reshape rather than a new subclass.

3. Façade. ChipAnalysis calls the kernel and builds the existing DressedResult.

Inspired by SuperGrad’s compute_energy_map: one-line jnp.argmax for greedy assignment, lax.scan argmax-with-masking for global greedy, and continuation along a parameter path. Generalized so references are pluggable as data and the path is a stacked eigvec tensor rather than a list of chip configurations.

Functions

assign_argmax(overlaps)	One-shot argmax per row.
assign_global_greedy(overlaps)	Global descending-overlap greedy via lax.scan with masking.
assign_rowwise_greedy(overlaps)	Confidence-ordered row-greedy assignment — O(n_labels * n_dressed).
compute_overlaps(reference, evecs)	|<ref_k | psi_j>|**2 matrix, shape (n_labels, n_dressed).
label_eigensystem(evecs, reference[, policy])	Pure-JAX labeling kernel.
phase_fixed_transform(labeling, evecs)	Phase-fixed transform from Kronecker basis to labeled dressed basis.
track_path(evals_path, evecs_path, ...[, policy])	Propagate a labeling along a parameter path of stacked eigensystems.

Classes

BareProductReference(dims)	Bare product basis states in Kronecker order.
EigenstateReference(vectors, keys)	Reference vectors are explicit eigenvectors with attached label keys.
Labeling(keys, indices, overlaps, margins, ...)	A labeled assignment of dressed eigenstates to user-meaningful keys.
LabelingPath(keys, indices, overlaps, ...)	Stacked labelings along a parameter path.

class quchip.chip.dressing.Labeling(keys, indices, overlaps, margins, duplicates)

Bases: object

A labeled assignment of dressed eigenstates to user-meaningful keys.

Array fields only — no embedded Reference object — so the dataclass is a clean pytree under jax.jit(), jax.vmap(), and jax.grad().

indices is the integer assignment (non-differentiable, conceptually stop_gradient’d). overlaps and margins are differentiable. Energies indexed via indices flow gradients through the gathered values: eigvals[labeling.indices[k]] is differentiable w.r.t. parameters that affect eigvals.

Parameters:

• keys (tuple[Any, ...])

• indices (Array)

• overlaps (Array)

• margins (Array)

• duplicates (Array)

keys: tuple[Any, ...]

indices: Array

overlaps: Array

margins: Array

duplicates: Array

class quchip.chip.dressing.LabelingPath(keys, indices, overlaps, margins, duplicates, swap_events)

Bases: object

Stacked labelings along a parameter path.

Each array field carries a leading n_steps axis vs. Labeling, plus a swap_events array surfacing where labels actually changed dressed-index between adjacent steps. A swap by itself is just a reassignment; combine it with a margin collapse (margins near zero) at the same step to identify avoided crossings vs. routine mode mixing.

Parameters:

• keys (tuple[Any, ...])

• indices (Array)

• overlaps (Array)

• margins (Array)

• duplicates (Array)

• swap_events (Array)

keys: tuple[Any, ...]

indices: Array

overlaps: Array

margins: Array

duplicates: Array

swap_events: Array

property final: Labeling

The labeling at the end of the path.

class quchip.chip.dressing.BareProductReference(dims)

Bases: object

Bare product basis states in Kronecker order.

Vectors are not materialized: bare products are the standard basis of the eigvec matrix in Kronecker order, so overlaps reduce to |evecs|**2 directly (see compute_overlaps()).

Parameters:

dims (tuple[int, ...])

dims: tuple[int, ...]

property keys: tuple[tuple[int, ...], ...]

class quchip.chip.dressing.EigenstateReference(vectors, keys)

Bases: object

Reference vectors are explicit eigenvectors with attached label keys.

Used for path continuation: the previous step’s selected eigvecs become the next step’s reference. vectors is row-major: each row is one reference state in the Kronecker basis.

Parameters:

• vectors (Array)

• keys (tuple[Any, ...])

vectors: Array

keys: tuple[Any, ...]

quchip.chip.dressing.compute_overlaps(reference, evecs)

|<ref_k | psi_j>|**2 matrix, shape (n_labels, n_dressed).

Specialized for BareProductReference to skip materializing an identity matrix: bare products are the Kronecker standard basis, so |<k|psi_j>|**2 = |evecs[k, j]|**2.

Parameters:

• reference (Any)

• evecs (Array)

Return type:

Array

quchip.chip.dressing.assign_argmax(overlaps)

One-shot argmax per row. Allows duplicate dressed-index assignments.

Returns (indices, chosen, margins). Cheap and jit-clean — matches SuperGrad’s default greedy mode. Use when the chip is in a weak-coupling regime where each bare label has a clearly-dominant dressed eigenstate.

Parameters:

overlaps (Array)

Return type:

tuple[Array, Array, Array]

quchip.chip.dressing.assign_global_greedy(overlaps)

Global descending-overlap greedy via lax.scan with masking.

At each step, finds the (label, dressed_idx) pair with the highest remaining overlap, assigns it, and masks out that label’s row and that column so neither is reused. Equivalent to the legacy Python sorted+set+dict path but pure JAX.

Margins are computed from the unmasked overlap matrix, so they describe per-label confidence independent of assignment order.

Requires n_labels <= n_dressed.

Parameters:

overlaps (Array)

Return type:

tuple[Array, Array, Array]

quchip.chip.dressing.assign_rowwise_greedy(overlaps)

Confidence-ordered row-greedy assignment — O(n_labels * n_dressed).

Labels (rows) are visited in descending row-max-overlap order; each takes its highest-overlap still-available dressed index, then that column is masked so the result is a permutation. Each scan step is a single masked-row argmax over n_dressed columns, so the whole assignment is O(n_labels * n_dressed) — for a square D = L**k overlap matrix that is O(L**(2k)) versus assign_global_greedy()’s O(L**(3k)) (a full L**(2k)-entry argmax at each of L**k steps). This is the SuperGrad / scqubits confidence-ordered greedy variant.

Semantics vs. global-greedy. Identical to assign_global_greedy() in the near-permutation regime (dispersive / weak hybridization), where each bare label has a clearly dominant dressed eigenstate. They can diverge on strongly-hybridized overlap matrices: global-greedy always commits the single largest remaining entry, whereas this visits rows in (original) row-max order. In exactly that regime the assignment is intrinsically ambiguous and Chip.dress already warns that bare labels are approximate. Pass policy=assign_global_greedy to label_eigensystem() for exact legacy semantics. Row-max ties are broken by ascending index (argsort).

Margins (top1 - top2 per label) come from the unmasked matrix via top_k, so they match assign_global_greedy()’s margins and avoid a full per-row sort. Requires n_labels <= n_dressed.

Parameters:

overlaps (Array)

Return type:

tuple[Array, Array, Array]

quchip.chip.dressing.label_eigensystem(evecs, reference, policy=<function assign_rowwise_greedy>)

Pure-JAX labeling kernel.

Parameters:

• evecs (jnp.ndarray) – (dim, n_dressed) eigenvectors as columns.

• reference (object) – Data-only object (BareProductReference or EigenstateReference) defining what label keys mean.

• policy (callable) – Assignment function overlaps -> (indices, chosen, margins). Default assign_rowwise_greedy() — the confidence-ordered row-greedy policy that Chip.dress runs (its sole live caller passes it explicitly).

Return type:

Labeling

Notes

Energies are looked up by the caller as evals[labeling.indices[k]]. policy is a static Python callable, resolved once at trace time, not a traced argument — under jax.jit it must be closed over or passed as a static_argnames entry, never as a dynamic (traced) argument.

quchip.chip.dressing.track_path(evals_path, evecs_path, initial_reference, policy=<function assign_global_greedy>)

Propagate a labeling along a parameter path of stacked eigensystems.

Step 0 is labeled against initial_reference. For each subsequent step, the reference becomes the previous step’s selected eigvecs (an EigenstateReference). This is SuperGrad’s continuation trick generalized: the path is a stacked eigvec tensor (typically the output of vmap(eigh) over a parameter grid), so it composes with any JAX-side parameter sweep.

Parameters:

• evals_path (jnp.ndarray) – (n_steps, n_dressed).

• evecs_path (jnp.ndarray) – (n_steps, dim, n_dressed).

• initial_reference (object) – Bootstrap reference at step 0.

• policy (callable) – Per-step assignment function. Defaults to assign_global_greedy().

Returns:

Stacked diagnostics including swap_events per step boundary.

Return type:

LabelingPath

quchip.chip.dressing.phase_fixed_transform(labeling, evecs)

Phase-fixed transform from Kronecker basis to labeled dressed basis.

Returns a complex (dim, n_labels) matrix whose columns are the dressed eigenstates assigned to labeling.keys, with each column phase-rotated so its largest-magnitude component is real-positive. Eigh’s default phase convention is gauge-arbitrary; this fix makes U a smooth function of physical parameters under vmap() and therefore safe for U.conj().T @ O @ U-style dressed-basis transforms (see ChipAnalysis.operator_in_dressed_basis()).

Only meaningful when labeling.indices is a permutation. If any labeling.duplicates is True, two columns of U will be identical — caller’s responsibility to guard.

Parameters:

• labeling (Labeling)

• evecs (Array)

Return type:

Array