Gradient with Zygote and Identity

[roflmaostc]

Hey,

I discovered a strange behaviour:

using Zygote
using Tullio

x = ones((5,5))

function f_working(arr)
    @tullio res = sqrt(arr[i, j])
    return res
end

function f_not_working(arr)
    @tullio res = identity(arr[i, j])
    return res
end

Calling the second function with Zygote:

julia> gradient(f_not_working, ones((3,3)))
ERROR: no gradient definition here!
Stacktrace:
 [1] (::Tullio.var"#217#218"{Tullio.Eval{var"#ℳ𝒶𝓀ℯ#9"{var"#𝒜𝒸𝓉!#8"},Nothing},Tuple{Array{Float64,2}},Array{Float64,1}})(::FillArrays.Fill{Float64,1,Tuple{Base.OneTo{Int64}}}) at <path>/.julia/packages/Tullio/bnAxO/src/grad/zygote.jl:7
 [2] (::Tullio.var"#195#back#219"{Tullio.var"#217#218"{Tullio.Eval{var"#ℳ𝒶𝓀ℯ#9"{var"#𝒜𝒸𝓉!#8"},Nothing},Tuple{Array{Float64,2}},Array{Float64,1}}})(::FillArrays.Fill{Float64,1,Tuple{Base.OneTo{Int64}}}) at <path>/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49
 [3] f_not_working at ./REPL[5]:3 [inlined]
 [4] (::typeof(∂(f_not_working)))(::Float64) at <path>/.julia/packages/Zygote/z3bQd/src/compiler/interface2.jl:0
 [5] (::Zygote.var"#37#38"{typeof(∂(f_not_working))})(::Float64) at <path>/.julia/packages/Zygote/z3bQd/src/compiler/interface.jl:45
 [6] gradient(::Function, ::Array{Float64,2}) at <path>/.julia/packages/Zygote/z3bQd/src/compiler/interface.jl:54
 [7] top-level scope at REPL[6]:1

Pure Zygote is able to make the derivative of identity, so it must be somehow due to the interplay of Tullio and Zygote.

Thanks,

Felix

Edit:
I think it also occurs in a more general case:

foo = sqrt

function f_not_working(arr)
    @tullio res = foo(arr[i, j])
    return res
end

[mcabbott]

Here's what's going on:

julia> @tullio res = identity(arr[i, j]) verbose=true 
┌ Warning: Symbolic gradient failed
│   err = no diffrule found for function identity(_).
└ @ Tullio ~/.julia/dev/Tullio/src/macro.jl:878
[ Info: left index ranges
┌ Info: reduction index ranges
│   i = Base.OneTo(5)
└   j = Base.OneTo(5)
25.0

That is, the default behaviour is to work out a symbolic gradient of the RHS, using Diffruls.jl. It turns out that nobody wrote a rule for identity, and certainly there will be no rule for foo. The macro sees only the symbol :foo, in general it has to assume that any function called :log behaves like Base.log, etc.

It is also possible to use dual numbers instead, in which case all that matters is that foo only contains operations which ForwardDiff understands:

using ForwardDiff
foo = sqrt
f_foo(arr) = @tullio res = foo(arr[i, j]) grad=Dual
gradient(f_foo, ones(5,5))[1]

[roflmaostc]

Hm, ok.
Yeah, what I wanted to try was to pass a function by reference to Tullio. So foo would have been
a differentiable function.

However, in my special case (similar to the metaprogramming #11 ) I could solve it with $foo where foo was assigned to be foo = :sqrt

Thanks!