No deriative rule found for `indexed_iterate`

[Roger-luo]

I think indexed_iterate should be defined either in CR or handled as intrinsic in AD engine?

MWE:

using Yota

A = Matrix(Hermitian(rand(4, 4)))

function test_f(A)
    λ, v = eigen(A)
    return sum(λ[1:1]) + sum(2.0 * v[:, 1])
end

Yota.grad(test_f, A)

julia> Yota.grad(test_f, A)
ERROR: No deriative rule found for op %13 = indexed_iterate(%4, 2, %11)::Tuple{Matrix{Float64}, Val{:done}} , try defining it using 

        ChainRulesCore.rrule(::typeof(Base.indexed_iterate), ::Eigen{Float64, Float64, Matrix{Float64}, Vector{Float64}}, ::Int64, ::Val{:vectors}) = ...

[dfdx]

Note that it complains about indexed_iterate(::Eigen, ...), which is indeed unknown to both Yota and ChainRules. For comparison, here's what happens when unpacking ordinary tuples:

julia> function test_f2(A)
           a, b = A[1], A[2]
           return a + b
       end
test_f2 (generic function with 1 method)

julia> Yota.grad(test_f2, A)
(0.5578878830250934, (ChainRulesCore.ZeroTangent(), [1.0 0.0 0.0 0.0; 1.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0]))

Rules for indexed_iterate() on tuples can be found here, but unfortunately I know little about Eigen type to immediately come up with the correct rule for it.

[Roger-luo]

I think it's the same as tuple since the indexed_iterate function is defined here: https://github.com/JuliaLang/julia/blob/master/base/tuple.jl#L88, I think there is no need to specialize it on tuple, since it is hardly (and shouldn't be) overloaded and more a compiler detail for the splitting syntax?

[dfdx]

It's a bit more complicated than that - Eigen doesn't directly overload indexed_iterate, but only iterate, which is called internally. Anyway, the output should be different (Tangent{Tuple} vs Tangent{Eigen}) and state type is also different (Int vs Val{field}). I'd say the rrule should be as follows, but it doesn't pass the test:

import Yota: YotaRuleConfig    # can be omitted
import Yota: gradcheck
import ChainRulesCore: rrule, NoTangent, ZeroTangent, Tangent

function rrule(::YotaRuleConfig, ::typeof(Base.indexed_iterate), t::Eigen, i::Int)
    y = Base.indexed_iterate(t, i)
    function indexed_iterate_pullback(dy)
        d_val, d_state = dy        
        return NoTangent(), Tangent{Eigen}(values=d_val), ZeroTangent()
    end
    return y, indexed_iterate_pullback
end

function rrule(::YotaRuleConfig, ::typeof(Base.indexed_iterate), t::Eigen, i::Int, state::Val{F}) where F
    y = Base.indexed_iterate(t, i, state)
    @assert F == :vectors
    function indexed_iterate_pullback(dy)
        d_val, d_state = dy        
        return NoTangent(), Tangent{Eigen}(vectors=d_val), ZeroTangent(), ZeroTangent()
    end
    return y, indexed_iterate_pullback
end

gradcheck(test_f, A)  # ==> false

I'm not very well educated in eigendecomposition, but I think the error might be related to numerical stability in eigen(...) / rrule(::typeof(eigen), ...). For example, if I change the definition of test_f to manual construction of Eigen, the test passes:

import Yota: __new__

function test_f(A)
    λ, v = __new__(Eigen{Float64, Float64, Matrix{Float64}, Vector{Float64}}, A[1, :], A)
    return sum(λ) + sum(v)
end

gradcheck(test_f, A)  # ==> true

[Roger-luo]

that's strange tho, __new__ should be the same as the constructor? Why indexed_iterate has something to do with numerical stability? the eigen rule passed rrule tests?

[dfdx]

that's strange tho, new should be the same as the constructor?

Not quite - __new__ creates the provided type without calling the constructor. In case of Eigen, if you write Eigen{...}(M, v), this inner method is invoked. But if you call __new__(Eigen{...}, M, v), Expr(:new, Eigen{...}, M, v) is used instead, which simply allocates the object and sets the fields accordingly.

This is needed to trace the creation of types with more complicated construction logic, e.g. when the default constructor (equivalent to Expr(:new, ...)) is hidden.

---

Why indexed_iterate has something to do with numerical stability?

It isn't, and that's the point. When using __new__(), we don't invoke eigen(), but simply construct the object and then call indexed_iterate() on it. The test passes in this case, which means that our rrule(::typeof(indexed_iterate), ::Eigen, ...) is probably correct.

On the other hand, if we check the original test_f() that additionally calls eigen(), the test fails. In fact, gradcheck fails even without further indexed_iterate:

gradcheck(A -> sum(eigen(A).vectors), A)  # ==> false

as well as test_rrule():

using ChainRulesTestUtils

test_rrule(eigen, A)  # ==> fails

I don't really think rrule(::typeof(eigen), ...) is implement incorrectly, more likely my test is numerically unstable. As I said earlier, I'm not good enough in eigendecomposition ¯_(ツ)_/¯