Training

Abstract base type for loss functions. Subtypes implement _loss_function.

L1 norm loss: minimizes sum(|T(x)|) to encourage sparsity.

L2Norm

L2 norm loss: sum(|T(x)|^2).

MSE loss with top-k truncation: ||x - T⁻¹(truncate(T(x), k))||². Field k is the number of kept coefficients.

_topk_mask(x::AbstractMatrix, k::Int) -> BitMatrix

Compute a boolean mask selecting the k elements of x with largest absolute value. Uses quickselect (O(n) average) to find the threshold, then a broadcast comparison.

Apply circuit to B images in a single einsum call. Returns (2,...,2,B) tensor.

Batched L1 loss: (1/B) * sum(|forward(images)|).

Batched L2 loss: (1/B) * sum(|forward(images)|^2).

Batched MSE loss: batched forward, per-image topk_truncate, batched inverse.

loss_function(tensors, m, n, optcode, pic::AbstractMatrix, loss; inverse_code=nothing)

Compute loss for a single image pic (2^m x 2^n) under the given circuit parameters.

loss_function(tensors, m, n, optcode, pics::Vector{<:AbstractMatrix}, loss; inverse_code=nothing, batched_optcode=nothing)

Average loss over a batch of images. Uses batched einsum if batched_optcode is provided.

Add a batch label to image input/output indices. Returns (batched_flat_code, batch_label).

Optimize contraction order for batched einsum. batch_size guides optimization but result works at any runtime batch size.

topk_truncate(x::AbstractMatrix, k::Integer)

Magnitude-based top-k truncation: keeps the k coefficients with largest absolute value, zeroing the rest. This is basis-agnostic — it does not assume any particular frequency layout.

Abstract base type for Riemannian manifolds used in circuit optimization.

U(1)^d manifold: each element is a unit complex number.

U(n) unitary group manifold. Tensors are n x n unitary matrices.

batched_adjoint(A::AbstractArray{T,3})

Batched conjugate transpose: C[:,:,k] = A[:,:,k]' for each slice k.

batched_inv(A::AbstractArray{T,3})

Batched matrix inverse: C[:,:,k] = inv(A[:,:,k]) for each slice k. Uses LU factorization for general matrices.

batched_matmul(A::AbstractArray{T,3}, B::AbstractArray{T,3})

Batched matrix multiply: C[:,:,k] = A[:,:,k] * B[:,:,k] for each slice k.

Classify a tensor as UnitaryManifold or PhaseManifold based on unitarity.

Group tensor indices by manifold type. Returns Dict{AbstractRiemannianManifold, Vector{Int}}.

is_unitary_general(t::AbstractMatrix)

Check if a square matrix satisfies U*U' ≈ I. Returns false for non-square matrices.

project(m::AbstractRiemannianManifold, points, euclidean_grads)

Project Euclidean gradients onto the tangent space at points. Batched over last dim.

Batched U(1)^d projection: im * imag(conj(z).*g) .* z.

Batched U(n) Riemannian projection: U * skew(U'G) on (d,d,n) arrays.

retract(m::AbstractRiemannianManifold, points, tangent_vec, α)

Retract from points along tangent_vec with step size α. Batched over last dim.

Batched U(1)^d retraction: normalize z + alpha*xi.

Batched Cayley retraction on U(n): (I - α/2·W)⁻¹(I + α/2·W)·U where W = Ξ·U'.

In-place version: pack into pre-allocated (d1, d2, n) array.

Pack selected matrices into a batched (d1, d2, n) array.

transport(m::AbstractRiemannianManifold, old_points, new_points, vec)

Parallel transport vec from old_points to new_points. Batched over last dim.

Batched U(1)^d parallel transport via re-projection.

Batched U(n) parallel transport via re-projection.

Unpack (d1, d2, n) array back into individual matrices.

Abstract base type for Riemannian optimizers.

Riemannian Adam optimizer (Becigneul & Ganea, 2019) with batched manifold operations.

Riemannian gradient descent with Armijo backtracking line search.

_batched_project(manifold_groups, point_batches, grad_buf_batches, euclidean_grads)

Batched Riemannian projection. Returns (rg_batches, grad_norm).

_compute_gradients(grad_fn, tensors)

Compute Euclidean gradients via grad_fn. Returns nothing on NaN/Inf.

optimize!(opt::RiemannianAdam, tensors, loss_fn, grad_fn; max_iter=100, tol=1e-6, loss_trace=nothing)

Run Riemannian Adam with momentum transport. Returns optimized tensors. When loss_trace::Vector{Float64} is provided, per-iteration losses are appended to it.

optimize!(opt::RiemannianGD, tensors, loss_fn, grad_fn; max_iter=100, tol=1e-6, loss_trace=nothing)

Run Riemannian gradient descent with Armijo line search. Returns optimized tensors. When loss_trace::Vector{Float64} is provided, per-iteration losses are appended to it.

Compute average loss over validation set.

Save a training checkpoint (basis + loss history).

Save training loss history to JSON.

Core training loop shared by all basis types. Returns (final_tensors, best_val_loss, train_losses, val_losses, step_train_losses). Uses optimize! from optimizers.jl for all optimization (GPU and CPU). Supports optimizers: RiemannianGD(), RiemannianAdam(), or symbols :gradient_descent, :adam.

Load training loss history from JSON. Returns a TrainingHistory object.

Save training history (from train_basis) to a JSON file.

Move array to CPU.

Move array to device. :gpu requires CUDA.jl via CUDAExt.

train_basis(::Type{EntangledQFTBasis}, dataset; m, n, entangle_phases, ...)

Train an EntangledQFTBasis on images. Same kwargs as QFTBasis plus entangle_phases.

Train a MERABasis on images. Same kwargs as TEBDBasis.

train_basis(::Type{QFTBasis}, dataset; m, n, loss, epochs, steps_per_image,
            optimizer, batch_size, device, ...)

Train a QFTBasis on images. Returns (basis, history). Key kwargs: optimizer (RiemannianGD()/RiemannianAdam()/:gradient_descent/:adam), batch_size, device (:cpu/:gpu).

train_basis(::Type{TEBDBasis}, dataset; m, n, phases, ...)

Train a TEBDBasis on images. Same kwargs as QFTBasis plus phases (initial TEBD gate phases).

Train QFTBasis from image files. Requires Images.jl.

_cache_key(flat_code, size_dict)

Compute a stable hash key for an einsum code + size dict combination.

optimize_code_cached(flat_code, size_dict, optimizer=TreeSA())

Like optimize_code(flat_code, size_dict, optimizer) but caches the result to disk. On cache hit, returns immediately without running the optimizer.

set_einsum_cache_dir!(dir::String)

Set the directory used for caching optimized einsum codes.