Skip to content

Expand TreeTN AD coverage with generic real/complex integration harness #375

@shinaoka

Description

@shinaoka

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 canonicalize and truncate
  • 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, and swap_site_indices are 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 f64 and Complex64
  • encode the scalar-specific convention differences explicitly:
    • forward-mode: standard JVP on C ~= R^2
    • reverse-mode: real-valued losses only for complex tests
  • 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_dense
  • canonicalize
  • truncate
  • add
  • evaluate
  • swap_site_indices
  • inner where 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_operator
  • contract_fit
  • contract_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 / Complex64 wrapper 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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions