Zygote: ERROR: LoadError: type NamedTuple has no field first

General Usage
1

Hi,
I’m trying to get started using Zygote in a numerical simulation. Basically I need to differentiate through a Laplace solver and am getting an error which I don’t know how to handle.

Here is the code:

using Plots
using Statistics
using LinearAlgebra
using Zygote


function invert_laplace(ρ,  Nz, Lz)
    # Solve Laplace equation
    # ρ: RHS function
    # Nz: Number of grid points
    # Lz: Length of the domain
    Δz = Lz / Nz
    zrg = range(0.0, step=Δz, length=Nz)
    # Squeeze a minus in here.
    invΔz² = -1.0 / Δz / Δz

    # Periodic boundary conditions. First row of A is just [1, 0, ... 0]
    A = diagm(0 => -2.0 * invΔz² * vcat([-0.5 / invΔz²], ones(Nz-1)), 
              1 => invΔz² * vcat([0.0], ones(Nz - 2)), 
              -1 => invΔz² * ones(Nz-1), -Nz+1 => [invΔz²])
    # Fix phi(x=0) = 0.
    ρvec = vcat([0.0], ρ(zrg[2:end]))
    ϕ_num = A \ ρvec
    return(ϕ_num)
end

function my_laplace_test(a)
    # Use a RHS that scales linear with and
    # solving ∂ₓ² ϕ = -a ρ is linear in a 
    # maximum(ϕ_num) should be a.
    Nz = 256
    Lz = 1.0
    ρ = z -> a * 4. * π * π * sin.(2π * z)
    ϕ_num = invert_laplace(ρ, Nz, Lz)
    return maximum(ϕ_num)
end

# Now let's try to differentiate through the solver:
res = gradient(my_laplace_test, 1.0)

But this gives me an error:

ERROR: LoadError: type NamedTuple has no field first

Stacktrace:

[1] getproperty(::NamedTuple{(:second,),Tuple{Array{Float64,1}}}, ::Symbol) at ./Base.jl:33
[2] macro expansion at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/lib/lib.jl:287 [inlined]
[3] (::Zygote.Jnew{Pair{Int64,Array{Float64,1}},Nothing,false})(::NamedTuple{(:second,),Tuple{Array{Float64,1}}}) at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/lib/lib.jl:279
[4] (::Zygote.var"#1727#back#165"{Zygote.Jnew{Pair{Int64,Array{Float64,1}},Nothing,false}})(::NamedTuple{(:second,),Tuple{Array{Float64,1}}}) at /Users/ralph/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59
[5] Pair at ./pair.jl:12 [inlined]
[6] Pair at ./pair.jl:15 [inlined]
[7] (::typeof(∂(Pair)))(::NamedTuple{(:second,),Tuple{Array{Float64,1}}}) at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/compiler/interface2.jl:0
[8] invert_laplace at /Users/ralph/source/julia/zygote_tests/test_laplace_direct.jl:25 [inlined]
[9] (::typeof(∂(invert_laplace)))(::Array{Float64,1}) at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/compiler/interface2.jl:0
[10] my_laplace_test at /Users/ralph/source/julia/zygote_tests/test_laplace_direct.jl:67 [inlined]
[11] (::Zygote.var"#41#42"{typeof(∂(my_laplace_test))})(::Float64) at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/compiler/interface.jl:40
[12] gradient(::Function, ::Float64) at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/compiler/interface.jl:49
[13] top-level scope at /Users/ralph/source/julia/zygote_tests/test_laplace_direct.jl:72
[14] include(::Function, ::Module, ::String) at ./Base.jl:380
[15] include(::Module, ::String) at ./Base.jl:368
[16] exec_options(::Base.JLOptions) at ./client.jl:296
[17] _start() at ./client.jl:506


Here [8] refers to the line where the matrix A is constructed. Does anyone have an idea what to do here?
2

Seems like it is hitting https://github.com/FluxML/Zygote.jl/pull/452/ which should be able to handle this.

Can we try to isolate the issue a bit? My thinking is that this shows up with the call to diagm maybe we are missing a constructor for that.

3

You’re right. If I replace construction of the A matrix with

A = Matrix(1.0I, Nz, Nz)

the code runs fine. As a quick fix I’ll probably remove all the syntax candy and construct A in a for-loop.

But should zygote in principle be able to handle this?

4

Constructing it like that might be slow; not sure

It should I think, pending a simple adjoint for the constructor

5

Here is the example that works:

using Statistics
using LinearAlgebra
using Zygote

# Test zygote's capability to differentiate through a laplace solver.
# This example initializes a profile 
# ρ = a 4π² sin(2 π x)
# and solves the Laplace equation
# ϕ = ∇⁻²ρ on x ∈ (0:1) with periodic boundary conditions:
# ϕ(x) = ϕ(x+1)
# Since the mean of ϕ is undetermined, we also fix ϕ(0) = 0.


function invert_laplace(ρ, Nz, Lz)
    # Solve Laplace equation: ϕ = ∇⁻² ρ on x ∈ (0:Lz).
    # ρ: RHS function
    # Nz: Number of grid points
    # Lz: Length of the domain
    # Returns: ϕ
    Δz = Lz / Nz
    # zrg = range(0.0, step=Δz, length=Nz)
    zrg = 0.0:Δz:Lz-Δz
    # Squeeze a minus in here.
    invΔz² = -1.0 / Δz / Δz

    # To fix ϕ(0) = 0 we have to get the equation
    # ϕ[0] * 1 = 0 in the linear system. To this by fixing A[1,1] = 1 and ρ[0] = 0.

    # # Periodic boundary conditions. First row of A is just [1, 0, ... 0]
    # A = diagm(0 => -2.0 * invΔz² * vcat([-0.5 / invΔz²], ones(Nz-1)), 
    #           1 => invΔz² * vcat([0.0], ones(Nz - 2)), 
    #           -1 => invΔz² * ones(Nz-1), -Nz+1 => [invΔz²])
    # # Fix phi(x=0) = 0.

    A0 = Zygote.Buffer(zeros(Nz, Nz))
    for n ∈ 1:Nz
        for m ∈ 1:Nz
            A0[n,m] = 0.0
        end
    end
    for n ∈ 2:Nz
        A0[n, n] = -2.0 * invΔz²
        A0[n, n-1] = invΔz²
        A0[n-1, n] = invΔz²
    end
    A0[1,1] = 1.0
    A0[1, 2] = 0.0
    A0[32,1] = invΔz²

    A = copy(A0)
    ρvec = vcat([0.0], ρ(zrg[2:end]))
    ϕ_num = A \ ρvec
    return(ϕ_num)
end

function my_laplace_test(a)
    # Use a RHS that scales linear with and solving ∂ₓ² ϕ = -a - ρ is linear in a 
    # maximum(ϕ_num) should be a.
    Nz = 256
    Lz = 1.0
    ρ = z -> a * 4. * π * π * sin.(2π * z)
    ϕ_num = invert_laplace(ρ, Nz, Lz)
    return maximum(ϕ_num)
end

# Now let's try to differentiate through the solver:
res = gradient(my_laplace_test, 1.0)

Now there are 2 issues:
1.

    zrg = range(0.0, step=Δz, length=Nz)
    zrg = 0.0:Δz:Lz-Δz

Using the first line gives an error:

  % /Applications/Julia-1.5.app/Contents/Resources/julia/bin/julia -i  test_laplace2.jl                                                                                                                                                                                                                                                     !10070
ERROR: LoadError: Need an adjoint for constructor StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}. Gradient is of type Array{Float64,1}
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] (::Zygote.Jnew{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},Nothing,false})(::Array{Float64,1}) at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/lib/lib.jl:301
 [3] (::Zygote.var"#1727#back#165"{Zygote.Jnew{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},Nothing,false}})(::Array{Float64,1}) at /Users/ralph/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59
 [4] StepRangeLen at ./range.jl:352 [inlined]
 [5] (::typeof(∂(StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}})))(::Array{Float64,1}) at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/compiler/interface2.jl:0
 [6] StepRangeLen at ./twiceprecision.jl:358 [inlined]
 [7] steprangelen_hp at ./twiceprecision.jl:344 [inlined]
 [8] (::typeof(∂(steprangelen_hp)))(::Array{Float64,1}) at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/compiler/interface2.jl:0
 [9] _range at ./twiceprecision.jl:447 [inlined]
 [10] (::typeof(∂(_range)))(::Array{Float64,1}) at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/compiler/interface2.jl:0
 [11] #range#43 at ./range.jl:91 [inlined]
 [12] (::typeof(∂(#range#43)))(::Array{Float64,1}) at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/compiler/interface2.jl:0 (repeats 2 times)
 [13] invert_laplace at /Users/ralph/source/julia/zygote_tests/test_laplace2.jl:12 [inlined]
 [14] (::typeof(∂(invert_laplace)))(::Array{Float64,1}) at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/compiler/interface2.jl:0
 [15] my_laplace_test at /Users/ralph/source/julia/zygote_tests/test_laplace2.jl:50 [inlined]
 [16] (::typeof(∂(my_laplace_test)))(::Float64) at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/compiler/interface2.jl:0
 [17] (::Zygote.var"#41#42"{typeof(∂(my_laplace_test))})(::Float64) at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/compiler/interface.jl:40
 [18] gradient(::Function, ::Float64) at /Users/ralph/.julia/packages/Zygote/7Jrhj/src/compiler/interface.jl:49
 [19] top-level scope at /Users/ralph/source/julia/zygote_tests/test_laplace2.jl:55
 [20] include(::Function, ::Module, ::String) at ./Base.jl:380
 [21] include(::Module, ::String) at ./Base.jl:368
 [22] exec_options(::Base.JLOptions) at ./client.jl:296
 [23] _start() at ./client.jl:506
in expression starting at /Users/ralph/source/julia/zygote_tests/test_laplace2.jl:55

Where would I need to define this custom adjoint constructor?

And 2.:
While the diagm syntax is nicer, the explicit construction is fine too. Of course it would be nice to have, but it’s not a must have. But just in case, what would need to be done to get custom adjoint constructors for diagm started?

6

It would be good to get a PR to Zygote or ChainRules with the adjoints for the specific constructors. Ditto for the first case. note that

julia> Flux.gradient(x -> sum(1:0.5:x), 5)
(nothing,)

does work fine.

You could also wrap it into a Zygote.ignore block