As someone relatively new to using ITensor, is there a straightforward way to integrate ITensor with an auto-differentiation library, such as Zygote.jl?
We have some support for differentiating functions involving ITensors and MPS/MPO with Zygote. The main limitation right now is that we haven’t added support for differentiating through tensor factorizations like SVD, QR, and eigendecompositions, which limits which code you can differentiate through. We get around that limitation to some extent for MPS/MPO operations by defining higher level derivative rules for things like MPO-MPS contraction and gate evolution (through the apply function), but the coverage is a little bit hit-and-miss and not documented. Feel free to try it out and let us know how it goes, the usefulness of the current support for automatic differentiation will greatly depend on your use case, patience to work around issues, etc.
As part of our ongoing rewrite of ITensors.jl, we have plans for lifting these restrictions and having much broader automatic differentiation (AD) support, which will be enabled by a few things:
- the code design in the rewrite is much simpler with much better abstractions and interfaces, which makes it easier for AD systems to work and also easier for us to write custom AD rules that are both simpler and provide better AD coverage,
- we are building our tensor factorizations on top of MatrixAlgebraKit.jl, which is a nicer interface for matrix factorization compared to the ones available in LinearAlgebra.jl and also already has, or very soon will have, full AD support, including for newer AD libraries like Enzyme.jl and Mooncake.jl, and
- we will likely directly target Mooncake.jl and Enzyme.jl instead of Zygote.jl as the AD libraries we will be compatible with. Mooncake.jl and Enzyme.jl have a big advantage over Zygote.jl in that they support mutation, which makes it a lot easier to have the AD system “just work” on a wider range of code, both ITensor library code and user code, so it really is much more automated (meaning fewer custom AD rules for library developers and more code that “just works” for users). Mooncake.jl has the added advantage that it works on type unstable code, which is particularly helpful for ITensors.jl since there is a baked-in level of type unstable code when performing tensor contractions (since the dimension of the output of a tensor contraction can only be known at runtime in general since we encode the tensor labels as runtime information attached to the tensors).
So we are very hopeful with the combination of the 3 points above that we can have full AD support for all ITensor operations, but that will only be available for general users once the rewrite is released, and unfortunately we don’t have a timeline for that (but maybe for basic operations we could have a pre-release soon).