[Yue-Zhengyuan]

In fixed-point gradient calculation, GeomSum and ManualIter previously called the same CTMRG pullback separately for its parameter p and environment x cotangents. Both solvers now accept a joint VJP and reuse both outputs, avoiding a redundant reverse pass per iteration:

$$\texttt{vjp}(\breve{v}) = (\text{ignored},\,\breve{v}\,\partial_p f,\,\breve{v}\,\partial_x f).$$

Update the fixedpoint_gradient docstring to describe the joint-VJP signature, and rename ∂f∂A to ∂f∂p to match the parameter p notation.

I find the notation dx, dy confusing, and will try to come up with more descriptive names as I learn about them.

[lkdvos]

actually looking at this a bit more, I wonder if we can also speed up the solvers by creating the following 3 functions to compute vjp's for:

f(state, env) = gauge_fixed_iteration(state, env) # both
g = Base.Fix1(f, state)
h = Base.Fix2(f, env)

It might be worth investigating if these vjp's are faster.

In any case, good catch, this does seem like it should speed up the gradients!

[Yue-Zhengyuan]

Can we do this in a follow up, so changes will be restricted to GeomSum and ManualIter?

[lkdvos]

left some minor comments about the docstrings but otherwise good to go for me, definitely also okay to put the other changes in a different PR.

[leburgel]

Up to a few very minor naming comments, looks perfect to me.

On the Fix1/Fix2 idea: I don’t think that will make a difference. Zygote differentiates through captured fields of a closure, so the state captured in Fix1(f, state) is still tracked, and TensorOperations will still compute (and force) the ket/bra cotangents on every pullback call. I first tried this approach in #430 hoping for a speedup, but it turned out that capturing the state in a closure gave identical allocations to passing it as an argument, which I was very annoyed about.

The only thing that helped was reading the local sandwich through a custom constant_site getindex helper whose gradient is explicitly ignored, so that the state cotangents are never tracked in the pure-environment VJP used in the linear solves.

I then gave up on doing the same for FixedPointGradient, since plumbing a similar hook all the way deep inside a CTMRG iteration was a total pain. Maybe with Enzyme the Fix1/Fix2 idea would work, but with our current setup it makes no difference.

[Yue-Zhengyuan]

New changes:

• The VJP is now collectively called ∂f.
• The argument order of iteration function is fixed to f(p, x) (parameter first).
• Some more renames from A to p for the variational parameters.
• Docstring of fixedpoint_gradient is overhauled to better explain the meaning of each argument.
• A TODO on trying Base.Fix idea once Enzyme is supported.

[Yue-Zhengyuan]

Benchmark

Measurements cover fixedpoint_gradient, including its VJP evaluations, and exclude forward CTMRG convergence and compilation. Comparing commit 4ed43fd with 50a272a.

Values below are upstream → current; percentage changes are relative to upstream, so negative values indicate reductions.

Heisenberg, D = 3, χ = 12

Solver	Time (ms)	Time change	Allocation (MiB)	Allocation change
Arnoldi	89.09 → 88.86	-0.26%	591.973 → 591.970	-0.0005%
GMRES	92.70 → 92.92	+0.23%	628.569 → 628.567	-0.0003%
BiCGStab	96.69 → 96.77	+0.08%	664.378 → 664.376	-0.0002%
GeomSum	187.36 → 96.03	-48.75%	1290.409 → 663.667	-48.5692%
ManualIter	188.08 → 96.72	-48.58%	1292.295 → 665.553	-48.4983%

[Yue-Zhengyuan]

@lkdvos Should we wrap alg_fixed, signs, env in ignore_derivatives? I did some tests; it has almost no effect on performance except a tiny reduction in allocations.