AssertionError in `back!`

[cscherrer]

I'm trying to get this code working:

using Tilde, Umlaut, Yota, ChainRulesCore, ChainRules

# From https://github.com/dfdx/Yota.jl/issues/113
# Guess this should really go into ChainRules
function ChainRulesCore.rrule(::typeof(tuple), args...)
    y = tuple(args...)
    return y, dy -> (NoTangent(), collect(dy)...)
end

m = @model begin
    p ~ Uniform()
    x ~ Bernoulli(p)
end

post = m() | (x = true,)
ℓ(p) = logdensityof(post, (p=p,))

# Works fine
ℓ(0.2)

# Throws an error
grad(ℓ, 0.2)

The error I get is

ERROR: AssertionError: zval isa Number || zval isa AbstractArray
Stacktrace:
  [1] back!(tape::Tape{Yota.GradCtx}; seed::Symbol)
    @ Yota ~/git/Yota.jl/src/grad.jl:200
  [2] gradtape!(tape::Tape{Yota.GradCtx}; seed::Symbol)
    @ Yota ~/git/Yota.jl/src/grad.jl:223
  [3] #gradtape#90
    @ ~/git/Yota.jl/src/grad.jl:245 [inlined]
  [4] make_rrule(f::Function, args::Accessors.PropertyLens{:ℓ})
    @ Yota ~/git/Yota.jl/src/cr_api.jl:128
  [5] rrule_via_ad(#unused#::Yota.YotaRuleConfig, f::Function, args::Accessors.PropertyLens{:ℓ})
    @ Yota ~/git/Yota.jl/src/cr_api.jl:170
  [6] rrule(#unused#::Yota.YotaRuleConfig, #unused#::typeof(Core._apply_iterate), #unused#::typeof(iterate), f::typeof(∘), args::Tuple{Accessors.PropertyLens{:ℓ}})
    @ Yota ~/git/Yota.jl/src/rulesets.jl:29
  [7] mkcall(::Function, ::Yota.YotaRuleConfig, ::Vararg{Any}; val::Missing, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Umlaut ~/git/Umlaut.jl/src/tape.jl:194
  [8] mkcall(::Function, ::Yota.YotaRuleConfig, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/tape.jl:179
  [9] record_primitive!(::Tape{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Yota ~/git/Yota.jl/src/grad.jl:49
 [10] record_or_recurse!(::Umlaut.Tracer{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:194
 [11] trace!(::Umlaut.Tracer{Yota.GradCtx}, ::Core.CodeInfo, ::Variable, ::Vararg{Variable})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:221
 [12] trace(f::Function, args::Accessors.PropertyLens{:ℓ}; ctx::Yota.GradCtx, fargtypes::Nothing, deprecated_kws::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:347
 [13] #gradtape#90
    @ ~/git/Yota.jl/src/grad.jl:244 [inlined]
 [14] make_rrule(f::Function, args::Accessors.PropertyLens{:ℓ})
    @ Yota ~/git/Yota.jl/src/cr_api.jl:128
 [15] rrule_via_ad(#unused#::Yota.YotaRuleConfig, f::Function, args::Accessors.PropertyLens{:ℓ})
    @ Yota ~/git/Yota.jl/src/cr_api.jl:170
 [16] rrule(#unused#::Yota.YotaRuleConfig, #unused#::typeof(Core._apply_iterate), #unused#::typeof(iterate), f::typeof(CompositionsBase.opcompose), args::Tuple{Accessors.PropertyLens{:ℓ}})
    @ Yota ~/git/Yota.jl/src/rulesets.jl:29
 [17] mkcall(::Function, ::Yota.YotaRuleConfig, ::Vararg{Any}; val::Missing, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Umlaut ~/git/Umlaut.jl/src/tape.jl:194
 [18] mkcall(::Function, ::Yota.YotaRuleConfig, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/tape.jl:179
 [19] record_primitive!(::Tape{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Yota ~/git/Yota.jl/src/grad.jl:49
 [20] record_or_recurse!(::Umlaut.Tracer{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:194
 [21] trace!(::Umlaut.Tracer{Yota.GradCtx}, ::Core.CodeInfo, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:221
 [22] record_or_recurse!(::Umlaut.Tracer{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:203
 [23] trace!(::Umlaut.Tracer{Yota.GradCtx}, ::Core.CodeInfo, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:221
 [24] record_or_recurse!(::Umlaut.Tracer{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:203
 [25] trace!(::Umlaut.Tracer{Yota.GradCtx}, ::Core.CodeInfo, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:221
 [26] record_or_recurse!(::Umlaut.Tracer{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:203
 [27] trace!(::Umlaut.Tracer{Yota.GradCtx}, ::Core.CodeInfo, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:221
 [28] record_or_recurse!(::Umlaut.Tracer{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:203
 [29] trace!(::Umlaut.Tracer{Yota.GradCtx}, ::Core.CodeInfo, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:221
 [30] record_or_recurse!(::Umlaut.Tracer{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:203
 [31] trace!(::Umlaut.Tracer{Yota.GradCtx}, ::Core.CodeInfo, ::Variable, ::Vararg{Variable})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:221
 [32] trace(f::Function, args::Float64; ctx::Yota.GradCtx, fargtypes::Nothing, deprecated_kws::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:347
 [33] #gradtape#90
    @ ~/git/Yota.jl/src/grad.jl:244 [inlined]
 [34] grad(f::typeof(ℓ), args::Float64; seed::Int64)
    @ Yota ~/git/Yota.jl/src/grad.jl:315
 [35] grad(f::typeof(ℓ), args::Float64)
    @ Yota ~/git/Yota.jl/src/grad.jl:307
 [36] top-level scope
    @ REPL[17]:1

By adding a @show in back!, I see the problem is that in this case, zval is Accessors.@optic _.ℓ, which has type Accessors.PropertyLens{:ℓ}.

Is this a bug? If not, could you give me some advice, for example how to know if this means I need another rrule? Teach a man to fish, etc :)

[dfdx]

I'm on foot right now, but can you try yo set seed argument to something of the same type?

[cscherrer]

I'm not sure if this is what you mean, but I tried changing the seed handling in back! to

    if (seed == 1) && (ndims(tape[z].val) > 0)
        error("Gradient of a vector-valued function requires a seed")
    elseif seed == :auto
        zval = tape[z].val
        # @assert zval isa Number || zval isa AbstractArray
        seed = setseed(zval)
    end

Then in my local Yota, I added

setseed(zval::Number) = one(zval)
setseed(zval::AbstractArray) = ones(eltype(zval), size(zval))

And in my test I added

Yota.setseed(lens::Accessors.PropertyLens) = lens

That got it past that point, but now I get an error in the getfield call here:

function Base.getproperty(op::AbstractOp, f::Symbol)
    if f == :typ
        return typeof(op.val)
    elseif f == :var
        return Variable(nothing, op)
    else
        getfield(op, f)
    end
end

Here's the stack trace:

ERROR: type Input has no field fn
Stacktrace:
  [1] getproperty(op::Input, f::Symbol)
    @ Umlaut ~/git/Umlaut.jl/src/tape.jl:98
  [2] todo_list(tape::Tape{Yota.GradCtx}, y::Variable)
    @ Yota ~/git/Yota.jl/src/grad.jl:132
  [3] todo_list
    @ ~/git/Yota.jl/src/grad.jl:131 [inlined]
  [4] back!(tape::Tape{Yota.GradCtx}; seed::Symbol)
    @ Yota ~/git/Yota.jl/src/grad.jl:210
  [5] gradtape!(tape::Tape{Yota.GradCtx}; seed::Symbol)
    @ Yota ~/git/Yota.jl/src/grad.jl:224
  [6] #gradtape#90
    @ ~/git/Yota.jl/src/grad.jl:245 [inlined]
  [7] make_rrule(f::Function, args::PropertyLens{:ℓ})
    @ Yota ~/git/Yota.jl/src/cr_api.jl:128
  [8] rrule_via_ad(#unused#::Yota.YotaRuleConfig, f::Function, args::PropertyLens{:ℓ})
    @ Yota ~/git/Yota.jl/src/cr_api.jl:170
  [9] rrule(#unused#::Yota.YotaRuleConfig, #unused#::typeof(Core._apply_iterate), #unused#::typeof(iterate), f::typeof(∘), args::Tuple{PropertyLens{:ℓ}})
    @ Yota ~/git/Yota.jl/src/rulesets.jl:29
 [10] mkcall(::Function, ::Yota.YotaRuleConfig, ::Vararg{Any}; val::Missing, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Umlaut ~/git/Umlaut.jl/src/tape.jl:194
 [11] mkcall(::Function, ::Yota.YotaRuleConfig, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/tape.jl:179
 [12] record_primitive!(::Tape{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Yota ~/git/Yota.jl/src/grad.jl:49
 [13] record_or_recurse!(::Umlaut.Tracer{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:194
 [14] trace!(::Umlaut.Tracer{Yota.GradCtx}, ::Core.CodeInfo, ::Variable, ::Vararg{Variable})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:221
 [15] trace(f::Function, args::PropertyLens{:ℓ}; ctx::Yota.GradCtx, fargtypes::Nothing, deprecated_kws::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:347
 [16] #gradtape#90
    @ ~/git/Yota.jl/src/grad.jl:244 [inlined]
 [17] make_rrule(f::Function, args::PropertyLens{:ℓ})
    @ Yota ~/git/Yota.jl/src/cr_api.jl:128
 [18] rrule_via_ad(#unused#::Yota.YotaRuleConfig, f::Function, args::PropertyLens{:ℓ})
    @ Yota ~/git/Yota.jl/src/cr_api.jl:170
 [19] rrule(#unused#::Yota.YotaRuleConfig, #unused#::typeof(Core._apply_iterate), #unused#::typeof(iterate), f::typeof(opcompose), args::Tuple{PropertyLens{:ℓ}})
    @ Yota ~/git/Yota.jl/src/rulesets.jl:29
 [20] mkcall(::Function, ::Yota.YotaRuleConfig, ::Vararg{Any}; val::Missing, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Umlaut ~/git/Umlaut.jl/src/tape.jl:194
 [21] mkcall(::Function, ::Yota.YotaRuleConfig, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/tape.jl:179
 [22] record_primitive!(::Tape{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Yota ~/git/Yota.jl/src/grad.jl:49
 [23] record_or_recurse!(::Umlaut.Tracer{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:194
 [24] trace!(::Umlaut.Tracer{Yota.GradCtx}, ::Core.CodeInfo, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:221
 [25] record_or_recurse!(::Umlaut.Tracer{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:203
 [26] trace!(::Umlaut.Tracer{Yota.GradCtx}, ::Core.CodeInfo, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:221
 [27] record_or_recurse!(::Umlaut.Tracer{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:203
 [28] trace!(::Umlaut.Tracer{Yota.GradCtx}, ::Core.CodeInfo, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:221
 [29] record_or_recurse!(::Umlaut.Tracer{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:203
 [30] trace!(::Umlaut.Tracer{Yota.GradCtx}, ::Core.CodeInfo, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:221
 [31] record_or_recurse!(::Umlaut.Tracer{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:203
 [32] trace!(::Umlaut.Tracer{Yota.GradCtx}, ::Core.CodeInfo, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:221
 [33] record_or_recurse!(::Umlaut.Tracer{Yota.GradCtx}, ::Function, ::Vararg{Any})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:203
 [34] trace!(::Umlaut.Tracer{Yota.GradCtx}, ::Core.CodeInfo, ::Variable, ::Vararg{Variable})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:221
 [35] trace(f::Function, args::Float64; ctx::Yota.GradCtx, fargtypes::Nothing, deprecated_kws::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Umlaut ~/git/Umlaut.jl/src/trace.jl:347
 [36] #gradtape#90
    @ ~/git/Yota.jl/src/grad.jl:244 [inlined]
 [37] grad(f::typeof(ℓ), args::Float64; seed::Int64)
    @ Yota ~/git/Yota.jl/src/grad.jl:316
 [38] grad(f::typeof(ℓ), args::Float64)
    @ Yota ~/git/Yota.jl/src/grad.jl:308
 [39] top-level scope
    @ REPL[29]:1

[dfdx]

Oh, sorry you had to dive into the seed implementation! I thought it would be enough to pass grad(...; seed=...), but it turns out the error comes from a dynamically generated rrule where seed is fixed to :auto. The same result could be achieved with:

    ...
    if (seed == 1) && (ndims(tape[z].val) > 0)
        error("Gradient of a vector-valued function requires a seed")
    elseif seed == :auto
        zval = tape[z].val
        # @assert zval isa Number || zval isa AbstractArray
        seed = zval isa AbstractArray ? ones(eltype(zval), size(zval)) : one(zval)
    end

...

Base.one(ℓ::Accessors.PropertyLens{:ℓ}) = ℓ    # or whatever would make a good seed

The second error also comes from a dynamically generated rrule. It seems like todo_list(), which should calculate the list of nodes to process, gets an Input op as the resulting node of a graph. From the stack trace I can't infer which function to blame, so I'll need some time to minimize the example and find the root cause.

[dfdx]

The problematic function is this:

∘(f) = f

It causes two complications for Yota:

1. It doesn't contain any function calls, so Yota doesn't understand what to do. I think I fixed it in my local code by adding identity at the end of such functions.
2. Given an optic, ∘() returns optic, while grad only knows how to automatically handle scalars and arrays. This is the same error we saw above, and the solution is to provide a good seed, i.e. a starting point for backpropagation that the very first pullback will receive.

Above I used a dummy implementation (one is used to produce a seed, but maybe we will need a dedicated function for this, idk yet):

Base.one(ℓ::Accessors.PropertyLens{:ℓ}) = ℓ 

Do you know a better implementation for this? I.e. what would be the unit value for an optic?

[cscherrer]

This is maybe naive, but I'd expect ∘(f) = f would just lead to any ∘(f) call ignoring ∘ and instead putting f itself on the tape. Where does this go wrong?

what would be the unit value for an optic?

My guess would be identity, but it's not clear to me what properties a seed is supposed to have.

Also, if the Accessors stuff is a problem for AD, I could probably bypass it by writing an rrule for models. I'm doing some refactoring now, but I should be able to try that within a couple of weeks if you think it would help.

[dfdx]

This is maybe naive, but I'd expect ∘(f) = f would just lead to any ∘(f) call ignoring ∘ and instead putting f itself on the tape. Where does this go wrong?

Somewhere down the road ∘(f...) is called instead. It lowers to Core._apply_iterate(iterate, ∘, (f,)), which is a higher-order function that leads to dynamic generation of a new rrule(∘, f), and this rule requires a custom seed. The good news is that I checked the uses of Core._apply_iterate(), and I now think we can be a bit smarter and trace ∘(f) to, well, just ∘(f), avoiding Core._apply_iterate() altogether. I'll need to touch both - the trace in Umlaut and the AD in Yota, but if you can wait for a couple of days, I'm really eager to get this case cleared.

[dfdx]

FYI: I'm working on removing Core._apply_iterate from the tape in the Umlaut#unapply-iterate-2 branch. The approach I took generally works, but a few tests still fail, plus I'll need some time to update Yota to the new interface. Please, bear with me!

[cscherrer]

Sounds good, thanks for the update!

[dfdx]

I have good news and bad news. Good news is that I've finally figured out ∘ and Core._apply_iterate cases. Or at least grad(ℓ, 0.2) doesn't throw an error anymore. Bad news is that instead of an error it produces:

(-1.6094379124341003, (ZeroTangent(), ZeroTangent()))

Assuming the model is actually differentiable from inputs to outputs, ZeroTangent() may be the result of a mistake in ChainRules, in my custom rrules or Yota itself. I also tried Zygote and Diffractor, but they both failed with compilation errors. It looks like the current model is a tough use case for AD systems.

One thing we can do though is to go through a number of examples using Tilde and detect where things start falling apart. Maybe you already have such a list to get started?

---

Here's the code I used:

using Tilde, Yota, Umlaut, ChainRulesCore, ChainRules

# THE FOLLOWING RRULES HAVE NOT BEEN TESTED! 
# DO NOT RELY ON THEM UNTIL PROPER TESTS ARE NOT IMPLEMENTED!

ChainRulesCore.rrule(::typeof(identity), x::Accessors.PropertyLens{:ℓ}) =
    (x, Δ -> (NoTangent(), Δ))


ChainRulesCore.rrule(::typeof(convert), ::Type{T}, x::T) where T =
    (x, Δ -> (NoTangent(), NoTangent(), Δ))


ChainRulesCore.rrule(::typeof(getfield), x, p::Symbol) =
    (getfield(x, p), Δ -> (NoTangent(), Tangent{typeof(x)}(; Dict(p=>Δ)...), NoTangent()))


function ChainRulesCore.rrule(::typeof(ifelse), cond, x, y)
    res = ifelse(cond, x, y)
    ifelse_pullback(Δ) = (NoTangent(), NoTangent(), (cond ? (Δ, ZeroTangent()) : (ZeroTangent(), Δ))...)
    return res, ifelse_pullback
end


function ChainRulesCore.rrule(::typeof(Core._apply_iterate), ::typeof(iterate), ::typeof(ifelse), args)
    res = ifelse(args...)
    cond = args[1]
    function splatted_ifelse_pullback(Δ)
        dargs = (
            NoTangent(),
            cond ? Δ : ZeroTangent(),
            cond ? ZeroTangent() : Δ
        )
        return NoTangent, NoTangent(), NoTangent(), dargs
    end
    return res, splatted_ifelse_pullback
end

    m = @model begin
        p ~ Uniform()
        x ~ Bernoulli(p)
    end

    post = m() | (x = true,)
    ℓ(p) = logdensityof(post, (p=p,))

    ℓ(0.2)

    grad(ℓ, 0.2)

and packages (Julia 1.7.3):

     Project Yota v0.7.4  (from the `main` branch)
      Status `~/work/Yota/Project.toml`
  [082447d4] ChainRules v1.39.2
  [d360d2e6] ChainRulesCore v1.15.3
  [26cc04aa] FiniteDifferences v0.12.24
  [872c559c] NNlib v0.8.8
  [92992a2b] Umlaut v0.4.1
  [73a6ac3c] Tilde v0.1.1 `https://github.com/cscherrer/Tilde.jl.git#b9d1a57`
  [e88e6eb3] Zygote v0.6.40
  [9f5e2b26] Diffractor v0.1.1 `https://github.com/JuliaDiff/Diffractor.jl.git#main`

[cscherrer]

This is great progress! I'm working on a big refactoring of Tilde now. I think it's getting close. I'd say it might be another week, though that might be too optimistic. You know how these things go :)

I think having some examples to test is a great idea. Maybe we can compare against ForwardDiff, or even FiniteDiff if we need to. Also, it's great to see Yota completing when Zygote and Diffractor fail, even if we have a bug to track down.

Thanks for your work on this!

[dfdx]

Indeed, both - ForwardDiff.derivative(ℓ, 0.2) and FiniteDiff.finite_difference_derivative(ℓ, 0.2) work well and return the same answer (5.0). I must admit at least for this specific example experience with ForwardDiff was very pleasant. Do you expect many models in Tilde to be low-dimensional and thus suitable e.g. for ForwardDiff?

Anyway, I'll be happy to provide support with differentiating for the next iteration of Tilde once it's ready.

[cscherrer]

Do you expect many models in Tilde to be low-dimensional and thus suitable e.g. for ForwardDiff?

I think the answer here is the same as for machine learning. Most examples and pedagogical models are low-dimensional, but for industrial use there are plenty of cases where we really need high dimensionality. To me, probprog actually subsumes machine learning, because you can always ask for just a point estimate. Also, log-densities from MeasureTheory can be treated as loss functions, and vice versa. Probprog models can also contain a deep network. I think the lines between these disciplines can become very blurred.

Anyway, I'll be happy to provide support with differentiating for the next iteration of Tilde once it's ready.

That would be great, thank you!