space: preserve tree structure in Jordan spectral results#58
Merged
Merged
Conversation
W3 — a tree's spectral results can now mirror the space treedef: - TreeSpace.spectrum(x, structured=True) returns a treedef-shaped pytree whose leaves hold each leaf's eigenvalue vector. The default flat spectrum() is unchanged (it is load-bearing for SpectralLpNormFunctional and the base trace/determinant reductions). - TreeSpectralDecomposition carries the producing space's treedef (static JAX pytree aux) and gains to_tree(), exposing the eigenvalues in the tree's own (possibly nested) structure. spectral_decompose tags it on. Round-trips decompose -> from_spectrum on nested trees. A structured spectrum is a treedef-matching pytree but NOT a space member (its leaves are eigenvalue- shaped), documented and covered by a check_member-rejects test. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
W3 — tree spectral results can now mirror the space
treedef, while the flat forms stay intact for the reduction/norm consumers.TreeSpace.spectrum(x, structured=True)returns atreedef-shaped pytree whose leaves hold each leaf's eigenvalue vector. The default flatspectrum(x)is unchanged — it is load-bearing forSpectralLpNormFunctional(Schatten norm over the whole direct-sum spectrum) and the basetrace/determinantreductions.TreeSpectralDecompositionnow carries the producing space'streedef(static JAX pytree aux) and gainsto_tree(), exposing the eigenvalues in the tree's own (possibly nested) structure.spectral_decomposetags it on.The
decompose → from_spectrumround-trip works on nested trees. A structured spectrum is atreedef-matching pytree but not a valid space member (its leaves are eigenvalue-shaped) — documented and enforced by acheck_member-rejects test.Notes
treedeffield defaults toNone(2-arg construction still works); flatspectrumoutput is byte-identical to before.SpectralLpNormFunctional's gradient over a tree domain was already unsupported (it 2-tuple-unpacksspectral_decompose); W3 doesn't change that — the regression test uses flat domains.Testing
tests/spaces/test_tree_spectral_decomposition.py: structured spectrum carriestreedef+flatten_treerecovers per-leaf spectra;to_tree(flat + nested); nesteddecompose→from_spectrumround-trip (incl. a Hermitian leaf and a complex-dtype tree); arity-1 branch; JAX pytree round-trip preservestreedef; structured spectrum rejected bycheck_member; and a flat-spectrum/SpectralLpNormFunctionalregression.Verification
An adversarial multi-agent review (JAX/pytree aux change, equality/hash, structure-preservation, consumer safety) found no bugs — the
PyTreeSpec-as-static-aux move round-trips throughjit/vmap(including the stacked-vmap-over-tree path), and the flat path is unchanged. Test gaps it surfaced (arity-1, non-member, complex) were added.🤖 Generated with Claude Code