Manifolds

The manifold geometry used by the optimizers: projections, retractions and batched linear algebra over stacked gate tensors.

Riemannian manifold abstraction with batched (d, d, n) operations.

Mirror of upstream src/manifolds.jl. Device-agnostic: works on CPU and GPU via JAX; no similar, copyto!, or mutable updates — pure functional.

Contents

AbstractRiemannianManifold	Interface the optimizers need from a manifold: project a Euclidean gradient onto the tangent space, retract from points along tangent by step alpha, and transport a tangent vector from old to new.
UnitaryManifold	U(d) unitary group manifold; tensors are d × d unitary matrices.
PhaseManifold	U(1)^d: each element is a unit complex number.
OrthogonalManifold	O(d) for real-valued d×d unitaries (subset of UnitaryManifold).
Unitary2qManifold	U(4) manifold for 2-qubit gates stored in (2, 2, 2, 2) tensor form.
Orthogonal2qManifold	O(4) for real-valued 2-qubit gates stored as (2, 2, 2, 2).
classify_manifold	Return the manifold appropriate to t based on its shape and unitarity.
group_by_manifold	{manifold: [indices]} bucket map.
stack_tensors	Pack selected matrices into a (d1, d2, n) batch.
unstack_tensors	Unpack a (*tensor_shape, n) batch back into a Python list, in place.
batched_matmul	C[:, :, k] = A[:, :, k] @ B[:, :, k] for each slice k.
batched_adjoint	C[:, :, k] = A[:, :, k].conj().T.
batched_inv	Batched matrix inverse via a transpose-and-invert trick.
is_unitary_2qubit	True for a (2, 2, 2, 2) tensor whose 4x4 reshape is unitary.
is_unitary_general	True if t @ t.conj().T ≈ I.