-
Notifications
You must be signed in to change notification settings - Fork 5
Expand TreeTN AD coverage with generic real/complex integration harness #375
Description
Summary
Broaden TreeTN automatic differentiation coverage beyond the initial #373 pass by introducing a generic test harness that runs the same integration logic for both f64 and Complex64, with forward-mode and reverse-mode conventions handled explicitly.
Current Status
Merged in #373 / commit 9334934:
TensorDynLen::backward()- TreeTN forward/backward AD integration tests for
to_dense - forward AD tangent preservation checks for
canonicalizeandtruncate - finite-difference validation for the current real-valued path
That closes the original gap tracked in #362, but the current coverage is still narrow:
- tests are mostly real-only
- there is no shared generic harness for
f64/Complex64 - representative TreeTN operations such as
add,evaluate, andswap_site_indicesare not covered by AD integration tests - higher-level operator/contraction paths remain untested from an AD perspective
Scope
Phase 1: Generic real/complex TreeTN AD harness
- add a shared test harness for TreeTN AD integration tests that can run for both
f64andComplex64 - encode the scalar-specific convention differences explicitly:
- forward-mode: standard JVP on
C ~= R^2 - reverse-mode: real-valued losses only for complex tests
- forward-mode: standard JVP on
- avoid duplicated test logic between real and complex variants
Phase 2: Representative TreeTN operation coverage
Add AD integration coverage for representative state-space operations:
to_densecanonicalizetruncateaddevaluateswap_site_indicesinnerwhere the chosen test shape makes the AD convention unambiguous
For each operation, prefer:
- forward-mode tangent-preservation or finite-difference checks
- reverse-mode gradient-propagation and finite-difference checks when the operation returns a tensor and the loss can be made real-valued
Phase 3: Follow-up algorithmic/operator coverage
Once the generic harness is in place, extend coverage to higher-level paths such as:
apply_linear_operatorcontract_fitcontract_zipup
This phase may be split again if any of these paths expose separate AD bugs.
Design Notes
Generic test structure
Follow the repo's generic-testing guidance from AGENTS.md:
- write test bodies once
- use thin
f64/Complex64wrapper tests - keep scalar-specific behavior in a small helper trait or helper functions
Complex convention
../chainrules-rs currently documents:
- forward-mode complex rules use standard JVP on
C ~= R^2 - reverse-mode complex rules use conjugate-Wirtinger for real-valued losses
TreeTN integration tests should match that convention.
Backward-mode losses for complex tests
Do not backpropagate through arbitrary complex-valued losses.
Use explicitly real-valued objectives such as:
real(inner(x, ref))- squared residuals
- norm-like objectives
Acceptance Criteria
- a generic real/complex AD test harness exists for TreeTN integration tests
- duplicated real/complex test logic is minimized
- representative TreeTN operations listed above are covered by integration tests where applicable
- complex tests use the documented forward/reverse AD conventions consistently
- any newly exposed AD bug is fixed or split into a dedicated follow-up issue with a reproducer