[Flux + Yao] Variable quantum circuits with trainable parameters

I’m doing quantum machine learning in Yao and Flux. The construction of the quantum circuits within Yao is fairly straightforward but contains a gate construction that depends on the number of qubits involved. When setting up the training loop, Flux/Zygote complains that it cannot differentiate the construction.

I have the quantum circuit set up as layers as

	    # Define the fixed structure for a single layer
	    function layer(α, β, γ)
	        circuit = chain(
	            repeat(Rx(α), 1:n),  # Global Rx(α)
	            repeat(Ry(β), 1:n)   # Global Ry(β)
	        )
	        zz_gates = [ZZGate(i, j, γ) for i in 1:n-1 for j in i+1:n]
	        return chain(n, circuit, chain(zz_gates...))
	    end
	
	    # Return a function that accepts parameters dynamically
	    return (params) -> begin
	        α_params, β_params, γ_params = params
	        layers = [layer(α_params[i], β_params[i], γ_params[i]) for i in 1:length(α_params)]
	        chain(layers...)
	    end
	end```

This creates three different gates, one layer of X rotations with a common parameter, a layer of Y rotations with a common parameter and ZZ gates as two-qubit gates, all with a common parameter.

When calculating the gradient with Flux as
```gs = Flux.gradient(() -> loss_fn_fixed(graph_state, label, model, params), Flux.Params([params]))```

It complains that `Mutating arrays is not supported -- called push!(Vector{LegibleLambdas.LegibleLambda{YaoBlocks.var"#87#88"{Tuple{LegibleLambdas.LegibleLambda{YaoBlocks.var"#120#121"{Int64, Pair{Int64, YaoBlocks.ConstGate.XGate}}}, LegibleLambdas.LegibleLambda{YaoBlocks.var"#132#133"{Pair{Int64, YaoBlocks.RotationGate{2, Float64, YaoBlocks.ConstGate.ZGate}}}}, LegibleLambdas.LegibleLambda{YaoBlocks.var"#120#121"{Int64, Pair{Int64, YaoBlocks.ConstGate.XGate}}}}}}}, ...)`
with error reference to the `zz_gates`.

The loss function is ```	### copied from Yao API out of desperation
	function _apply_fallback!(r::AbstractArrayReg{D,T}, b::AbstractBlock) where {D,T}
	    r.state = mat(T, b) * r.state
		return r
	end

	
	function loss_fn_fixed(graph_state, label, model, params)
	    # Apply the model (quantum circuit with parameters) to the graph state
	    state = _apply_fallback!(graph_state, model(params))
	    
	    # Compute the expectation value of Z^{\otimes n}
	    prediction = expectation_z_tensor_n(state)
	    
	    # Mean squared error between prediction and true label
	    return (prediction - label)^2
	end```
where I had to copy the apply_fallback and modify it since it also directly mutates the state register `r`.

I have already constructed the `fixed_quantum_circuit` circuit to separate construction and feeding of parameters in an attempt to resolve the problem, the prior attempt ran into similar issues of trying to differentiate the model construction.
Is there a straightforward way to fix the model architecture such that Flux/Zygote does not try to differentiate the model?

You might be looking for Zygote.@ignore?

Otherwise it’s hard to guess without at least the full stacktrace (to know within what function Zygote sees array mutation) and preferably code which includes everything (i.e. can be pasted into a new session).