# Automatic Differentiation of quantum master equation using Zygote

**URL:** <https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715>\
**Category:** Specific Domains\
**Tags:** zygote, autodiff, ordinarydiffeq\
**Created:** [April 4, 2025, 4:08pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715 "2025-04-04T16:08:46Z")\
**Posts on this page:** 15\
**Page:** 1

<div class="post-metadata">

**Author:** ![albertomercurio](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albertomercurio/32/27051_2.png) [@albertomercurio](https://discourse.julialang.org/u/albertomercurio)\
**Post date:** [April 4, 2025, 4:08pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/1 "2025-04-04T16:08:46Z")

</div>

Hello,

I need to differentiate a quantum master equation over a parameter, using QuantumToolbox.jl. Under the hood it is simply an Ordinary Differential Equation involving SciMLOperators.

I’m able to differentiate the Schrödinger equation with the following code

```julia-auto
using QuantumToolbox
using OrdinaryDiffEq
using SciMLSensitivity
using Zygote
using Enzyme
using BenchmarkTools

##

const N = 16
F = 1
γ = 1

const a = destroy(N)

const ψ0 = fock(N, 0)

# coef1(p, t) = p.Δ
coef1(p, t) = p[1]

const H = QobjEvo(a' * a, coef1) + F * (a' + a)
const c_ops = [sqrt(γ) * a]

const L = liouvillian(H, c_ops)

function my_f_sesolve(p)
  tlist = range(0, 1.5, 100)

  sol = sesolve(H, ψ0, tlist, progress_bar=Val(false), params=p, saveat=[tlist[end]], sensealg = BacksolveAdjoint(autojacvec=EnzymeVJP()))

  # return real(abs2(sol.states[end].data[end]))
  return real(expect(a' * a, sol.states[end]))
end

Δ = 1.0
params = [Δ]
my_f_sesolve(params) # 1.8585255970284336

##

Zygote.gradient(my_f_sesolve, params) # ([-2.7374779226913737],)

```

The internal code is probably too complicated to be explained here, but the working principle is basic linear algebra between matrices and vectors.

Current limitations of autodiff for this situation can be found in [this issue](https://github.com/SciML/SciMLSensitivity.jl/issues/1139), especially in the last comment where I summarize everything.

However, if I try to differentiate the master equation it fails. This case should be equivalent to the previous one, where the SciMLOperator involved in the ODE is only bigger than the Schrödinger case. Here the code

```julia-auto
function my_f_mesolve(p)
  tlist = range(0, 1.5, 100)

  sol = mesolve(H, ψ0, tlist, c_ops, progress_bar=Val(false), params=p, saveat=[tlist[end]], sensealg = BacksolveAdjoint(autojacvec=EnzymeVJP()))

  # return real(abs2(sol.states[end].data[end]))
  return real(expect(a' * a, sol.states[end]))
end

##

Δ = 1.0
params = [Δ]
my_f_mesolve(params) # 0.925041907046927

##

Zygote.gradient(my_f_mesolve, params)

```

And the Error:

> **Summary**
>
> ```julia-auto
> ERROR: Need an adjoint for constructor QuantumObject{OperatorQuantumObject, Dimensions{1, Tuple{Space}}, Matrix{ComplexF64}}. Gradient is of type Diagonal{Float64, Vector{Float64}}
> Stacktrace:
> [1] error(s::String)
> @ Base ./error.jl:35
> [2] (::Zygote.Jnew{QuantumObject{…}, Nothing, false})(Δ::Diagonal{Float64, Vector{…}})
> @ Zygote ~/.julia/packages/Zygote/HdT4O/src/lib/lib.jl:324
> [3] (::Zygote.var"#2230#back#327"{Zygote.Jnew{QuantumObject{…}, Nothing, false}})(Δ::Diagonal{Float64, Vector{Float64}})
> @ Zygote ~/.julia/packages/ZygoteRules/CkVIK/src/adjoint.jl:72
> [4] QuantumObject
> @ ~/.julia/packages/QuantumToolbox/YmsAS/src/qobj/quantum_object.jl:59 [inlined]
> [5] (::Zygote.Pullback{Tuple{…}, Any})(Δ::Diagonal{Float64, Vector{…}})
> @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface2.jl:0
> [6] *
> @ ~/.julia/packages/QuantumToolbox/YmsAS/src/qobj/arithmetic_and_attributes.jl:72 [inlined]
> [7] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Diagonal{Float64, Vector{…}})
> @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface2.jl:0
> [8] expect
> @ ~/.julia/packages/QuantumToolbox/YmsAS/src/qobj/functions.jl:58 [inlined]
> [9] my_f_mesolve
> @ ~/GitHub/Research/Undef/Autodiff QuantumToolbox/autodiff.jl:65 [inlined]
> [10] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64)
> @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface2.jl:0
> [11] (::Zygote.var"#88#89"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Float64)
> @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface.jl:97
> [12] gradient(f::Function, args::Vector{Float64})
> @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface.jl:154
> [13] top-level scope
> @ ~/GitHub/Research/Undef/Autodiff QuantumToolbox/autodiff.jl:76
> Some type information was truncated. Use `show(err)` to see complete types.
> 
> ```

The differentiation of the ODE itself seem to work, since if I return `real(abs2(sol.states[end].data[end]))` instead of `real(expect(a' * a, sol.states[end]))` it works. So, apparently the `expect` function doesn’t work in the case of `mesolve`, but it works in the case of `sesolve`.

To give more context, a `QuantumObject` like the `a` variable here is a `struct` containing three fields

- The `data` field: the actual matrix or vector representing the state or the operator
- The `type` and `dimensions` fields: some structs to store the information of the quantum system

## Conclusions

It seems that I have to define some adjoint for the `QuantumObject` struct, whatever it means. If so, I don’t understand why I don’t need it for the Schrödinger equation case, but I need it for the master equation.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [April 4, 2025, 7:24pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/2 "2025-04-04T19:24:13Z")

</div>

> [@albertomercurio](#):
>
> Current limitations of autodiff for this situation can be found in [this issue](https://github.com/SciML/SciMLSensitivity.jl/issues/1139), especially in the last comment where I summarize everything.
> 
> However, if I try to differentiate the master equation it fails. This case should be equivalent to the previous one, where the SciMLOperator involved in the ODE is only bigger than the Schrödinger case. Here the code

Yes I need to get to this… I know, it’s still in my email.

---

<div class="post-metadata">

**Author:** ![albertomercurio](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albertomercurio/32/27051_2.png) [@albertomercurio](https://discourse.julialang.org/u/albertomercurio)\
**Post date:** [April 4, 2025, 8:41pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/3 "2025-04-04T20:41:54Z")

</div>

Don’t worry, in this thread I’m asking something different from the issue on GitHub. Here I get the error

```julia
ERROR: Need an adjoint for constructor QuantumObject{OperatorQuantumObject, Dimensions{1, Tuple{Space}}, Matrix{ComplexF64}}. Gradient is of type Diagonal{Float64, Vector{Float64}}

```

When I return `real(expect(a' * a, sol.states[end]))` rather than a simple matrix element `real(abs2(sol.states[end].data[end]))`.

So it is something apparently not related to the ODE itself, but to my custom type.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [April 4, 2025, 8:42pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/4 "2025-04-04T20:42:45Z")

</div>

> [@albertomercurio](#):
>
> `QuantumObject`

Oh you’re missing an adjoint for this constructor function.

---

<div class="post-metadata">

**Author:** ![albertomercurio](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albertomercurio/32/27051_2.png) [@albertomercurio](https://discourse.julialang.org/u/albertomercurio)\
**Post date:** [April 4, 2025, 8:45pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/5 "2025-04-04T20:45:33Z")

</div>

Ok, but how is it possible that with `sesolve` (the first ODE) it doesn’t need it? It only happens if I call the `expect` function, which is basically

```julia
expect(A::QuantumObject, B::QuantumObject) = tr(A.data * B.data)

```

where the data field is a simple matrix.

How can I fix it?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [April 4, 2025, 8:48pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/6 "2025-04-04T20:48:31Z")

</div>

> [@albertomercurio](#):
>
> where the data field is a simple matrix.
> 
> How can I fix it?

Because the generic adjoint of getfield uses the constructor adjoint in the backpass. You just need something like this:

> <https://github.com/SciML/SciMLBase.jl/blob/v2.82.0/ext/SciMLBaseZygoteExt.jl#L77-L80>

I.e., the derivative of `getfield` is basically just an identity, so the return is `A.data`, but then for the backwards pass if `J = I` then `J' = I`, so the derivative is just the identity, so you just build a new `QuantumObject` with the pullback value.

---

<div class="post-metadata">

**Author:** ![albertomercurio](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albertomercurio/32/27051_2.png) [@albertomercurio](https://discourse.julialang.org/u/albertomercurio)\
**Post date:** [April 5, 2025, 9:19am UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/7 "2025-04-05T09:19:37Z")

</div>

I tried first this

```julia
Zygote.@adjoint function Zygote.literal_getproperty(A::QuantumObject, ::Val{:data})
  A.data, p̄ -> (QuantumObject(p̄, A.type, A.dimensions),)
end

```

but I still get the same error.

I then tried

```julia
Zygote.@adjoint function QuantumObject(A::AbstractArray, type, dimensions)
  QuantumObject(A, type, dimensions), p̄ -> (QuantumObject(p̄, type, dimensions),)
end

```

But the I get the error

```julia
BoundsError: attempt to access Tuple{Nothing, QuantumObject{OperatorQuantumObject, Dimensions{1, Tuple{Space}}, Diagonal{Float64, Vector{Float64}}}} at index [3]

```

I really don’t know why it works for `sesolve` but not for `mesolve`. They are both ODE with SciMLOperators, where in the second case the operator is just slightly different but still of the same form.

The error seems to be related to the last line and not to the ODE itself.

---

<div class="post-metadata">

**Author:** ![albertomercurio](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albertomercurio/32/27051_2.png) [@albertomercurio](https://discourse.julialang.org/u/albertomercurio)\
**Post date:** [April 5, 2025, 12:30pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/8 "2025-04-05T12:30:31Z")

</div>

I tried a simpler function to test Zygote.jl on `QuantumObjects`, and it works without defining custom adjoints for the `QuantumObject` constructor.

```julia
const ψgg = tensor(basis(2, 1), basis(2, 1))
const ψee = tensor(basis(2, 0), basis(2, 0))

function my_fidelity(p)
  ψ = normalize(p[1] * ψgg + p[2] * ψee)

  return fidelity(ψ, ψgg)
end

p = [0.3, 0.4]
my_fidelity(p) # 0.6

Zygote.gradient(my_fidelity, p) # ([1.28, -0.96],)

```

Here `ψgg` and `ψee` are `QuantumObject` types containing a `Vector` as data. The `fidelity` function just computes `dot(ψ.data, ψgg.data)`.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [April 5, 2025, 12:33pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/9 "2025-04-05T12:33:22Z")

</div>

> [@albertomercurio](#):
>
> ```julia
> BoundsError: attempt to access Tuple{Nothing, QuantumObject{OperatorQuantumObject, Dimensions{1, Tuple{Space}}, Diagonal{Float64, Vector{Float64}}}} at index [3]
> 
> ```

You’re missing the derivative w.r.t. the other parts. I’d normally do that one with ChainRules though. But:

```julia
Zygote.@adjoint function QuantumObject(A::AbstractArray, type, dimensions)
  QuantumObject(A, type, dimensions), p̄ -> (QuantumObject(p̄, type, dimensions), ChainRules.ZeroTangent(), ChainRules.ZeroTangent())
end

```

since the derivative w.r.t. the other two parts is structurally zero since they are discrete.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [April 5, 2025, 12:33pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/10 "2025-04-05T12:33:54Z")

</div>

> [@albertomercurio](#):
>
> I tried a simpler function to test [Zygote.jl](https://juliaregistries.github.io/General/packages/redirect_to_repo/Zygote) on `QuantumObjects`, and it works without defining custom adjoints for the `QuantumObject` constructor.
> 
> ```julia
> 
> ```

It depends on how the values are accessed.

---

<div class="post-metadata">

**Author:** ![albertomercurio](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albertomercurio/32/27051_2.png) [@albertomercurio](https://discourse.julialang.org/u/albertomercurio)\
**Post date:** [April 5, 2025, 1:18pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/11 "2025-04-05T13:18:11Z")

</div>

Ok, fair, thanks.

I tried and now the error is

```julia-auto
ERROR: MethodError: no method matching *(::QuantumObject{OperatorQuantumObject, Dimensions{…}, Diagonal{…}}, ::Adjoint{ComplexF64, Matrix{…}})

```

It is weird, since it is trying to do the product between a QuantumObject and a matrix.

However, the `sesolve` case no longer works with this adjoint definition, and it gives

```julia-auto
ERROR: MethodError: no method matching QuantumObject(::@NamedTuple{data::SparseMatrixCSC{ComplexF64, Int64}, type::Nothing, dimensions::Nothing}, ::OperatorQuantumObject, ::Dimensions{1, Tuple{Space}})

```

I tried to change the function to something like

```julia-auto
Zygote.@adjoint function QuantumObject(A, type, dimensions)
  back(p̄) = (QuantumObject(p̄, type, dimensions), ChainRules.ZeroTangent(), ChainRules.ZeroTangent())
  back(p̄::NamedTuple{(:data, )}) = (QuantumObject(p̄.data, type, dimensions), ChainRules.ZeroTangent(), ChainRules.ZeroTangent())
  QuantumObject(A, type, dimensions), back
end

```

But it is still calling the first method of the `back` function.

### EDIT

By changing the definition to

```julia-auto
Zygote.@adjoint function QuantumObject(A, type, dimensions)
  back(p̄) = (QuantumObject(p̄, type, dimensions), ChainRules.ZeroTangent(), ChainRules.ZeroTangent())
  back(p̄::NamedTuple{(:data, :type, :dimensions)}) = (QuantumObject(p̄.data, type, dimensions), ChainRules.ZeroTangent(), ChainRules.ZeroTangent())
  QuantumObject(A, type, dimensions), back
end

```

now `sesolve` has the same error as for `mesolve`.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [April 5, 2025, 1:24pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/12 "2025-04-05T13:24:52Z")

</div>

> [@albertomercurio](#):
>
> `back(p̄::NamedTuple{(:data, :type, :dimensions)})`

back(p̄::NamedTuple{(:data, :type, :dimensions)}) that’s a missing projection / means that the constructor adjoint wasn’t called.

---

<div class="post-metadata">

**Author:** ![albertomercurio](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albertomercurio/32/27051_2.png) [@albertomercurio](https://discourse.julialang.org/u/albertomercurio)\
**Post date:** [April 5, 2025, 1:41pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/13 "2025-04-05T13:41:04Z")

</div>

Why this is happening? How to fix it?

### EDIT

I tried to play a bit. I found the following case working

```julia-auto
Zygote.@adjoint function QuantumObject(data, type, dims)
  obj = QuantumObject(data, type, dims)
  function pullback(Δobj)
      # Return gradient for data, and ignore type and dims
      return (Δobj.data, nothing, nothing)
  end
  function pullback(Δobj::AbstractArray)
    # Return gradient for data, and ignore type and dims
    return (Δobj, nothing, nothing)
  end
  return obj, pullback
end

```

I need the definition of the `pullback` also for standard `AbstractArray`s.

Is it reasonable? If yes, how can I convert it to a ChainRulesCore.jl `rrule`?

---

<div class="post-metadata">

**Author:** ![albertomercurio](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albertomercurio/32/27051_2.png) [@albertomercurio](https://discourse.julialang.org/u/albertomercurio)\
**Post date:** [April 5, 2025, 2:45pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/14 "2025-04-05T14:45:15Z")

</div>

In the direction of defining the adjoint with ChainRulesCore.jl, I tried this implementation

```julia
function ChainRulesCore.rrule(::Type{QuantumObject}, data, type, dimensions)
  obj = QuantumObject(data, type, dimensions)
  f_pullback(Δobj) = (NoTangent(), Δobj.data, NoTangent(), NoTangent())
  f_pullback(Δobj_data::AbstractArray) = (NoTangent(), Δobj_data, NoTangent(), NoTangent())
  return obj, f_pullback
end

```

It seems to work. I had to change `typeof(QuantumObject)` with `Type{QuantumObject}`, maybe because it is a constructor? Also, the pullback function is not using the `Tangent` type. But I didn’t managed to use it.

With this implementation, my code is working also with ChainRulesCore.jl.

Do you think that these two definitions (Zygote and ChainRulesCore) are correct and reasonable?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [April 5, 2025, 7:14pm UTC](https://discourse.julialang.org/t/automatic-differentiation-of-quantum-master-equation-using-zygote/127715/15 "2025-04-05T19:14:43Z")

</div>

> [@albertomercurio](#):
>
> Do you think that these two definitions (Zygote and ChainRulesCore) are correct and reasonable?

yes