Taking the gradient of `x->sum(cumsum(x))` triggers segmentation fault

[YingboMa]

MWE:

julia> versioninfo()
Julia Version 1.0.0
Commit 5d4eaca0c9 (2018-08-08 20:58 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: Intel(R) Core(TM) i7-6820HQ CPU @ 2.70GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-6.0.0 (ORCJIT, skylake)
Environment:
  JULIA_PKG3_PRECOMPILE = 1

julia> using Zygote: gradient
[ Info: Precompiling Zygote [e88e6eb3-aa80-5325-afca-941959d7151f]
┌ Warning: error processing module Zygote signature expressions Expr[:(Tuple{Core.Typeof(typeinf_code3), Method, Any, SimpleVector, Bool, Params})] from :(@eval Core.Compiler function typeinf_code3(method::Method, @nospecialize(atypes), sparams::SimpleVector, run_optimizer::Bool, params::Params)
│           code = code_for_method(method, atypes, sparams, params.world)
│           code === nothing && return (nothing, Any)
│           ccall(:jl_typeinf_begin, Cvoid, ())
│           result = InferenceResult(code)
│           frame = InferenceState(result, false, params)
│           frame === nothing && return (nothing, Any)
│           if typeinf(frame) && run_optimizer
│               opt = OptimizationState(frame)
│               optimize(opt, result.result)
│               (opt.src).inferred = true
│           end
│           ccall(:jl_typeinf_end, Cvoid, ())
│           frame.inferred || return (nothing, Any)
│           return frame
│       end)
└ @ Revise ~/.julia/packages/Revise/51oQc/src/exprutils.jl:94

julia> gradient(x->sum(cumsum(x)), zeros(5))

signal (11): Segmentation fault
in expression starting at no file:0
CreateConstInBoundsGEP1_32 at /buildworker/worker/package_linux64/build/usr/include/llvm/IR/IRBuilder.h:1292 [inlined]
emit_sparam at /buildworker/worker/package_linux64/build/src/codegen.cpp:3262
emit_expr at /buildworker/worker/package_linux64/build/src/codegen.cpp:3923
emit_ssaval_assign at /buildworker/worker/package_linux64/build/src/codegen.cpp:3615
emit_stmtpos at /buildworker/worker/package_linux64/build/src/codegen.cpp:3801 [inlined]
emit_function at /buildworker/worker/package_linux64/build/src/codegen.cpp:6254
jl_compile_linfo at /buildworker/worker/package_linux64/build/src/codegen.cpp:1159
jl_fptr_trampoline at /buildworker/worker/package_linux64/build/src/gf.c:1794
jl_apply_generic at /buildworker/worker/package_linux64/build/src/gf.c:2182
cumsum at ./accumulate.jl:116 [inlined]
_forward at /home/scheme/.julia/packages/Zygote/g8WMA/src/compiler/interface2.jl:0
#3 at ./REPL[3]:1 [inlined]
_forward at /home/scheme/.julia/packages/Zygote/g8WMA/src/compiler/interface2.jl:0
jl_fptr_trampoline at /buildworker/worker/package_linux64/build/src/gf.c:1829
jl_apply_generic at /buildworker/worker/package_linux64/build/src/gf.c:2182
_forward at /home/scheme/.julia/packages/Zygote/g8WMA/src/compiler/interface.jl:21
forward at /home/scheme/.julia/packages/Zygote/g8WMA/src/compiler/interface.jl:27
gradient at /home/scheme/.julia/packages/Zygote/g8WMA/src/compiler/interface.jl:32
jl_fptr_trampoline at /buildworker/worker/package_linux64/build/src/gf.c:1829
jl_apply_generic at /buildworker/worker/package_linux64/build/src/gf.c:2182
do_call at /buildworker/worker/package_linux64/build/src/interpreter.c:324
eval_value at /buildworker/worker/package_linux64/build/src/interpreter.c:428
eval_stmt_value at /buildworker/worker/package_linux64/build/src/interpreter.c:363 [inlined]
eval_body at /buildworker/worker/package_linux64/build/src/interpreter.c:686
jl_interpret_toplevel_thunk_callback at /buildworker/worker/package_linux64/build/src/interpreter.c:799
unknown function (ip: 0xfffffffffffffffe)
unknown function (ip: 0x7fc57b60581f)
unknown function (ip: 0x7)
jl_interpret_toplevel_thunk at /buildworker/worker/package_linux64/build/src/interpreter.c:808
jl_toplevel_eval_flex at /buildworker/worker/package_linux64/build/src/toplevel.c:787
jl_toplevel_eval_in at /buildworker/worker/package_linux64/build/src/builtins.c:622
eval at ./boot.jl:319
jl_apply_generic at /buildworker/worker/package_linux64/build/src/gf.c:2182
eval_user_input at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.0/REPL/src/REPL.jl:85
run_backend at /home/scheme/.julia/packages/Revise/51oQc/src/Revise.jl:766
#58 at ./task.jl:259
jl_fptr_trampoline at /buildworker/worker/package_linux64/build/src/gf.c:1829
jl_apply_generic at /buildworker/worker/package_linux64/build/src/gf.c:2182
jl_apply at /buildworker/worker/package_linux64/build/src/julia.h:1536 [inlined]
start_task at /buildworker/worker/package_linux64/build/src/task.c:268
unknown function (ip: 0xffffffffffffffff)
Allocations: 35859323 (Pool: 35853151; Big: 6172); GC: 79
[1]    22985 segmentation fault (core dumped)  $HOME/julia-1.0.0/bin/julia

[chriselrod]

I get

julia> versioninfo()
Julia Version 1.0.0
Commit dadc85b246* (2018-08-09 18:04 UTC)
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: AMD Ryzen Threadripper 1950X 16-Core Processor
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-6.0.0 (ORCJIT, znver1)

julia> using Zygote: gradient
[ Info: Precompiling Zygote [e88e6eb3-aa80-5325-afca-941959d7151f]
[ Info: Precompiling IRTools [7869d1d1-7146-5819-86e3-90919afe41df]

julia> gradient(x->sum(cumsum(x)), zeros(5))
ERROR: AssertionError: not implemented
Stacktrace:
 [1] record_branches!(::Core.Compiler.IRCode) at /home/chris/.julia/packages/Zygote/UUjpC/src/compiler/reverse.jl:65
 [2] #Primal#39(::Nothing, ::Type, ::Core.Compiler.IRCode) at /home/chris/.julia/packages/Zygote/UUjpC/src/compiler/reverse.jl:179
 [3] Type at ./none:0 [inlined]
 [4] #Adjoint#65 at /home/chris/.julia/packages/Zygote/UUjpC/src/compiler/reverse.jl:377 [inlined]
 [5] (::getfield(Core, Symbol("#kw#Type")))(::NamedTuple{(:varargs,),Tuple{Nothing}}, ::Type{Zygote.Adjoint}, ::Core.Compiler.IRCode) at ./none:0
 [6] _lookup_grad(::Type) at /home/chris/.julia/packages/Zygote/UUjpC/src/compiler/emit.jl:121
 [7] #s18#608 at /home/chris/.julia/packages/Zygote/UUjpC/src/compiler/interface2.jl:17 [inlined]
 [8] #s18#608(::Any, ::Any, ::Any) at ./none:0
 [9] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any,N} where N) at ./boot.jl:506
 [10] _cumsum! at ./accumulate.jl:55 [inlined]
 [11] (::Zygote.J{Tuple{typeof(Base._cumsum!),Array{Float64,1},Array{Float64,1},Int64,Base.ArithmeticRounds},Tuple{typeof(Base._cumsum!),Array{Float64,1},Array{Float64,1},Int64,Base.ArithmeticRounds,getfield(Zygote, Symbol("##82#92")),Array{Any,1},Array{Any,1},Array{Bool,1}}})(::Array{Float64,1}) at /home/chris/.julia/packages/Zygote/UUjpC/src/compiler/interface2.jl:0
 [12] #cumsum!#571 at ./accumulate.jl:51 [inlined]
 [13] (::Zygote.J{Tuple{getfield(Base, Symbol("##cumsum!#571")),Int64,typeof(cumsum!),Array{Float64,1},Array{Float64,1}},Tuple{getfield(Base, Symbol("##cumsum!#571")),Int64,typeof(cumsum!),Array{Float64,1},Array{Float64,1},Zygote.J{Tuple{typeof(Base._cumsum!),Array{Float64,1},Array{Float64,1},Int64,Base.ArithmeticRounds},Tuple{typeof(Base._cumsum!),Array{Float64,1},Array{Float64,1},Int64,Base.ArithmeticRounds,getfield(Zygote, Symbol("##82#92")),Array{Any,1},Array{Any,1},Array{Bool,1}}},Zygote.J{Tuple{Type{Base.ArithmeticStyle},Type{Float64}},Tuple{}},Zygote.J{Tuple{typeof(eltype),Array{Float64,1}},Tuple{typeof(eltype),Array{Float64,1},Zygote.J{Tuple{typeof(eltype),Type{Array{Float64,1}}},Tuple{typeof(eltype)}},getfield(Zygote, Symbol("##79#89"))}}}})(::Array{Float64,1}) at /home/chris/.julia/packages/Zygote/UUjpC/src/compiler/interface2.jl:0
 [14] #cumsum! at ./none:0 [inlined]
 [15] (::Zygote.J{Tuple{getfield(Base, Symbol("#kw##cumsum!")),NamedTuple{(:dims,),Tuple{Int64}},typeof(cumsum!),Array{Float64,1},Array{Float64,1}},Tuple{getfield(Base, Symbol("#kw##cumsum!")),NamedTuple{(:dims,),Tuple{Int64}},typeof(cumsum!),Array{Float64,1},Array{Float64,1},Zygote.J{Tuple{typeof(haskey),NamedTuple{(:dims,),Tuple{Int64}},Symbol},Tuple{typeof(haskey),NamedTuple{(:dims,),Tuple{Int64}},Symbol,Zygote.J{Tuple{typeof(isdefined),NamedTuple{(:dims,),Tuple{Int64}},Symbol},Tuple{typeof(isdefined)}}}},Array{Any,1},Array{Any,1},Array{Any,1},Array{Any,1},Array{Any,1},Array{Bool,1},Array{Any,1},Array{Bool,1}}})(::Array{Float64,1}) at /home/chris/.julia/packages/Zygote/UUjpC/src/compiler/interface2.jl:0
 [16] (::Zygote.J{Tuple{getfield(Base, Symbol("##cumsum#572")),Int64,typeof(cumsum),Array{Float64,1}},Tuple{getfield(Base, Symbol("##cumsum#572")),Int64,typeof(cumsum),Array{Float64,1},Zygote.J{Tuple{getfield(Base, Symbol("#kw##cumsum!")),NamedTuple{(:dims,),Tuple{Int64}},typeof(cumsum!),Array{Float64,1},Array{Float64,1}},Tuple{getfield(Base, Symbol("#kw##cumsum!")),NamedTuple{(:dims,),Tuple{Int64}},typeof(cumsum!),Array{Float64,1},Array{Float64,1},Zygote.J{Tuple{typeof(haskey),NamedTuple{(:dims,),Tuple{Int64}},Symbol},Tuple{typeof(haskey),NamedTuple{(:dims,),Tuple{Int64}},Symbol,Zygote.J{Tuple{typeof(isdefined),NamedTuple{(:dims,),Tuple{Int64}},Symbol},Tuple{typeof(isdefined)}}}},Array{Any,1},Array{Any,1},Array{Any,1},Array{Any,1},Array{Any,1},Array{Bool,1},Array{Any,1},Array{Bool,1}}},Zygote.J{Tuple{Type{NamedTuple{(:dims,),T} where T<:Tuple},Tuple{Int64}},Tuple{UnionAll,Tuple{Int64},Zygote.J{Tuple{Type{NamedTuple{(:dims,),Tuple{Int64}}},Tuple{Int64}},Tuple{DataType,Tuple{Int64},getfield(Zygote, Symbol("##188#back2#140")){Zygote.Jnew{NamedTuple{(:dims,),Tuple{Int64}},Nothing}},getfield(Zygote, Symbol("##160#back2#120")){getfield(Zygote, Symbol("##116#118")){1,Int64}}}}}},getfield(Zygote, Symbol("##148#back2#110")){typeof(identity)},Zygote.J{Tuple{typeof(similar),Array{Float64,1},Type{Float64}},Tuple{typeof(similar),Array{Float64,1},DataType,getfield(Zygote, Symbol("##888#back2#495")){getfield(Zygote, Symbol("##493#494"))},getfield(Zygote, Symbol("##496#501"))}},Zygote.J{Tuple{typeof(Base.promote_op),typeof(Base.add_sum),Type{Float64},Type{Float64}},Tuple{typeof(Base.promote_op)}}}})(::Array{Float64,1}) at ./accumulate.jl:91
 [17] #cumsum at ./none:0 [inlined]
 [18] (::Zygote.J{Tuple{getfield(Base, Symbol("#kw##cumsum")),NamedTuple{(:dims,),Tuple{Int64}},typeof(cumsum),Array{Float64,1}},Tuple{getfield(Base, Symbol("#kw##cumsum")),NamedTuple{(:dims,),Tuple{Int64}},typeof(cumsum),Array{Float64,1},Zygote.J{Tuple{typeof(haskey),NamedTuple{(:dims,),Tuple{Int64}},Symbol},Tuple{typeof(haskey),NamedTuple{(:dims,),Tuple{Int64}},Symbol,Zygote.J{Tuple{typeof(isdefined),NamedTuple{(:dims,),Tuple{Int64}},Symbol},Tuple{typeof(isdefined)}}}},Array{Any,1},Array{Any,1},Array{Any,1},Array{Any,1},Array{Any,1},Array{Any,1},Array{Bool,1},Array{Any,1},Array{Bool,1}}})(::Array{Float64,1}) at /home/chris/.julia/packages/Zygote/UUjpC/src/compiler/interface2.jl:0
 [19] cumsum at ./accumulate.jl:116 [inlined]
 [20] (::Zygote.J{Tuple{typeof(cumsum),Array{Float64,1}},Tuple{typeof(cumsum),Array{Float64,1},Zygote.J{Tuple{getfield(Base, Symbol("#kw##cumsum")),NamedTuple{(:dims,),Tuple{Int64}},typeof(cumsum),Array{Float64,1}},Tuple{getfield(Base, Symbol("#kw##cumsum")),NamedTuple{(:dims,),Tuple{Int64}},typeof(cumsum),Array{Float64,1},Zygote.J{Tuple{typeof(haskey),NamedTuple{(:dims,),Tuple{Int64}},Symbol},Tuple{typeof(haskey),NamedTuple{(:dims,),Tuple{Int64}},Symbol,Zygote.J{Tuple{typeof(isdefined),NamedTuple{(:dims,),Tuple{Int64}},Symbol},Tuple{typeof(isdefined)}}}},Array{Any,1},Array{Any,1},Array{Any,1},Array{Any,1},Array{Any,1},Array{Any,1},Array{Bool,1},Array{Any,1},Array{Bool,1}}},Zygote.J{Tuple{Type{NamedTuple{(:dims,),T} where T<:Tuple},Tuple{Int64}},Tuple{UnionAll,Tuple{Int64},Zygote.J{Tuple{Type{NamedTuple{(:dims,),Tuple{Int64}}},Tuple{Int64}},Tuple{DataType,Tuple{Int64},getfield(Zygote, Symbol("##188#back2#140")){Zygote.Jnew{NamedTuple{(:dims,),Tuple{Int64}},Nothing}},getfield(Zygote, Symbol("##160#back2#120")){getfield(Zygote, Symbol("##116#118")){1,Int64}}}}}},getfield(Zygote, Symbol("##148#back2#110")){typeof(identity)}}})(::Array{Float64,1}) at /home/chris/.julia/packages/Zygote/UUjpC/src/compiler/interface2.jl:0
 [21] #7 at ./REPL[4]:1 [inlined]
 [22] (::getfield(Zygote, Symbol("##66#67")){Zygote.J{Tuple{getfield(Main, Symbol("##7#8")),Array{Float64,1}},Tuple{getfield(Main, Symbol("##7#8")),Array{Float64,1},getfield(Zygote, Symbol("##540#543")){getfield(Zygote, Symbol("##957#back2#542")){getfield(Zygote, Symbol("##539#541")){Array{Float64,1}}}},Zygote.J{Tuple{typeof(cumsum),Array{Float64,1}},Tuple{typeof(cumsum),Array{Float64,1},Zygote.J{Tuple{getfield(Base, Symbol("#kw##cumsum")),NamedTuple{(:dims,),Tuple{Int64}},typeof(cumsum),Array{Float64,1}},Tuple{getfield(Base, Symbol("#kw##cumsum")),NamedTuple{(:dims,),Tuple{Int64}},typeof(cumsum),Array{Float64,1},Zygote.J{Tuple{typeof(haskey),NamedTuple{(:dims,),Tuple{Int64}},Symbol},Tuple{typeof(haskey),NamedTuple{(:dims,),Tuple{Int64}},Symbol,Zygote.J{Tuple{typeof(isdefined),NamedTuple{(:dims,),Tuple{Int64}},Symbol},Tuple{typeof(isdefined)}}}},Array{Any,1},Array{Any,1},Array{Any,1},Array{Any,1},Array{Any,1},Array{Any,1},Array{Bool,1},Array{Any,1},Array{Bool,1}}},Zygote.J{Tuple{Type{NamedTuple{(:dims,),T} where T<:Tuple},Tuple{Int64}},Tuple{UnionAll,Tuple{Int64},Zygote.J{Tuple{Type{NamedTuple{(:dims,),Tuple{Int64}}},Tuple{Int64}},Tuple{DataType,Tuple{Int64},getfield(Zygote, Symbol("##188#back2#140")){Zygote.Jnew{NamedTuple{(:dims,),Tuple{Int64}},Nothing}},getfield(Zygote, Symbol("##160#back2#120")){getfield(Zygote, Symbol("##116#118")){1,Int64}}}}}},getfield(Zygote, Symbol("##148#back2#110")){typeof(identity)}}}}}})(::Int64) at /home/chris/.julia/packages/Zygote/UUjpC/src/compiler/interface.jl:28
 [23] gradient(::Function, ::Array{Float64,1}) at /home/chris/.julia/packages/Zygote/UUjpC/src/compiler/interface.jl:34
 [24] top-level scope at none:0

Which is interesting, because we both seem to be on Julia v1 and Zygote commit e88e6eb3.

[MikeInnes]

Probably best to fix this by adding a gradient definition for cumsum (perhaps easier to add to Flux and I can port it).

I think it should just be a case of flipping the summed dimensions of the sensitivity and cumsuming those, e.g.

julia> Tracker.ngradient(x -> sum(cumsum(x)), randn(5))
([5.0, 4.0, 3.0, 2.0, 1.0],)

[pshashk]

@YingboMa My attempt at implementing cumsum gradient for Flux is here: FluxML/Flux.jl#388 . Not sure why gradtest fails though.

[MikeInnes]

This now gives an appropriate error:

julia> gradient(x -> sum(cumsum(x)), [1, 2, 3])
ERROR: Mutating arrays is not supported