Too long waiting for the zygote AD for the first execution during the optimization

[Confusio]

When the env dimension chi is large, it costs too long time (more than 30 mins when χbond = 4
χenv = 64 for the j1-j2 model in the example) waiting for the AD to start. Monitoring the cpu, I find that only one cpu-core is running on the multi-cores node during the waiting. How to work around the long waiting time?

[lkdvos]

Given how heavily we rely on Zygote for the AD, this is not something we have a lot of influence over. A reasonable strategy is to precompile this such that you only pay this cost once, see https://quantumkithub.github.io/PEPSKit.jl/dev/man/precompilation/

[Confusio]

Thanks! Looks like this strategy isn’t much better than just grabbing a coffee and waiting :( .

[lkdvos]

Except you would only have to go and grab a coffee once, instead of every time you restart your Julia session.

[Confusio]

Just wondering—any chance to switch to a different AD library (eg. Enzyme.jl) someday (of course not near future), given the mutation issue with Zygote?

[lkdvos]

This would indeed be a very nice upgrade, both for overall performance, but also for convenience, considering the mutation etc.
What's kind of holding this back, apart from the rather large amount of work involved, is the clear choice of what engine would be best.
Enzyme seems to be the engine of choice for smaller pieces of code, where performance is key.
The drawback being that the error messages are even more cryptic, and I don't actually know how well it performs on scales such as ours.
On the other hand, MoonCake.jl seems a very promising choice, but they still indicate that they are in the development stage.
Neither of these systems is compatible with the other, so writing the required custom rules to make everything work needs to be more or less duplicated, and given our quite large stack of code, I'm somewhat hesitant to get this started.

All of this to say, I would absolutely love to switch, I'm just kind of scared of the amount of work and time this would take, and I don't want to have to switch and then end up having to switch again.
Realistically though, it would be feasible to start smaller, for example checking if we could switch just the contractions to Enzyme while keeping all other parts to Zygote, and seeing if that makes a noticeable difference.
Help and opinions of course always welcome