Zygote struggles with circuits that have a nested structure such as put => chain => put

[dandanua]

Hi!

I have noticed that Zygote struggles with gradient computation for circuits that have nested put components - e.g., chain(10, [put(10, (i,j) => chain(2, [put(2, 1=>Rx(0)), put(2, 2=>Ry(0))])) for i=1:5, j=6:10]). The gradient computation time (it could be the compilation time, though) grows enormously in comparison to the equivalent "flat" circuit without nested puts.

I don't know what's the issue, but I think it would be good to have a method in YaoBlocks.Optimise that removes nested puts.

[GiggleLiu]

There is a method called Optimise.canonicalize:

julia> c = chain(10, [put(10, (i,j) => chain(2, [put(2, 1=>Rx(0)), put(2, 2=>Ry(0))])) for i=1:5, j=6:10]);

julia> Optimise.canonicalize(c)
nqubits: 10
chain
├─ put on (1)
│  └─ rot(X, 0.0)
├─ put on (6)
│  └─ rot(Y, 0.0)
├─ put on (2)
│  └─ rot(X, 0.0)
...

I guess this is what you want. I think it is worth mentioning in the documentation.

[dandanua]

Yes, this is what I wanted, it should help avoid that problem with Zygote. I haven't found it in the docs, it's definitely worth mentioning.