Differentiating across `add_edges` causes Zygote to segfault

Reproduce el fallo con el ejemplo de Julia y Zygote del issue y, después, inspecciona add_edges en GNNGraphs/src/transform.jl:326, donde entra el stack trace. Compara el comportamiento con el issue relacionado de Zygote. Se considera terminado cuando diferenciar loss_fn devuelve un gradiente sin provocar un segfault ni finalizar el proceso de Julia.

Using the add_edges function on a path that is differentiated with Zygote.jl causes a segfault and kills the julia process. The version is v1.1.0
Could be related to https://github.com/FluxML/Zygote.jl/issues/1580

using GraphNeuralNetworks
using Zygote

s = [1, 2]
t = [2, 3]
w = Float32[0.5, 0.7]

g = GNNGraph((s, t, w))

function loss_fn(wnew)
    # fixed new edges, only weights are differentiable
    snew = [1, 2]
    tnew = [2, 3]

    g2 = add_edges(g, (snew, tnew, wnew))

    return sum(get_edge_weight(g2))
end

wnew = Float32[0.1, 0.2]

gradient(loss_fn, wnew)

The resulting segfault is

[205339] signal 11 (1): Segmentation fault
in expression starting at REPL[8]:2
emit_unboxed_coercion at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/intrinsics.cpp:394 [inlined]
emit_unbox at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/intrinsics.cpp:458
emit_condition at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/codegen.cpp:6089
emit_condition at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/codegen.cpp:6102 [inlined]
emit_function at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/codegen.cpp:9294
jl_emit_code at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/codegen.cpp:9750
jl_emit_codeinst at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/codegen.cpp:9821
jl_emit_codeinst_to_jit_impl at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/jitlayers.cpp:770
jl_add_codeinst_to_jit at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/gf.c:3376
add_codeinsts_to_jit! at ./../usr/share/julia/Compiler/src/typeinfer.jl:1433
typeinf_ext_toplevel at ./../usr/share/julia/Compiler/src/typeinfer.jl:1440 [inlined]
typeinf_ext_toplevel at ./../usr/share/julia/Compiler/src/typeinfer.jl:1448
jfptr_typeinf_ext_toplevel_80206.1 at /home/numerik/amartine/.julia/juliaup/julia-1.12.6+0.x64.linux.gnu/lib/julia/sys.so (unknown line)
jl_apply at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/julia.h:2391 [inlined]
jl_type_infer at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/gf.c:466
jl_compile_method_internal at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/gf.c:3516
_jl_invoke at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/gf.c:4006 [inlined]
ijl_apply_generic at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/gf.c:4214
#add_edges#170 at /scratch/htc/amartine/julia_thunderstorm/packages/GNNGraphs/Z9SPx/src/transform.jl:326 [inlined]
_pullback at /scratch/htc/amartine/julia_thunderstorm/packages/Zygote/55SqB/src/compiler/interface2.jl:0
_pullback at /scratch/htc/amartine/julia_thunderstorm/packages/Zygote/55SqB/src/compiler/interface2.jl:81
unknown function (ip: 0x7f90f3be47fb) at (unknown file)
loss_fn at ./REPL[6]:5 [inlined]
_pullback at /scratch/htc/amartine/julia_thunderstorm/packages/Zygote/55SqB/src/compiler/interface2.jl:0
pullback at /scratch/htc/amartine/julia_thunderstorm/packages/Zygote/55SqB/src/compiler/interface.jl:96
pullback at /scratch/htc/amartine/julia_thunderstorm/packages/Zygote/55SqB/src/compiler/interface.jl:94 [inlined]
gradient at /scratch/htc/amartine/julia_thunderstorm/packages/Zygote/55SqB/src/compiler/interface.jl:153
unknown function (ip: 0x7f90f3b2b6c6) at (unknown file)
jl_apply at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/julia.h:2391 [inlined]
do_call at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/interpreter.c:123
eval_value at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/interpreter.c:243
eval_stmt_value at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/interpreter.c:194 [inlined]
eval_body at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/interpreter.c:707
jl_interpret_toplevel_thunk at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/interpreter.c:898
jl_toplevel_eval_flex at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/toplevel.c:1035
__repl_entry_eval_expanded_with_loc at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/usr/share/julia/stdlib/v1.12/REPL/src/REPL.jl:301
jl_apply at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/julia.h:2391 [inlined]
jl_f_invokelatest at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/builtins.c:881
toplevel_eval_with_hooks at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/usr/share/julia/stdlib/v1.12/REPL/src/REPL.jl:308
toplevel_eval_with_hooks at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/usr/share/julia/stdlib/v1.12/REPL/src/REPL.jl:312
toplevel_eval_with_hooks at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/usr/share/julia/stdlib/v1.12/REPL/src/REPL.jl:305 [inlined]
eval_user_input at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/usr/share/julia/stdlib/v1.12/REPL/src/REPL.jl:330
repl_backend_loop at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/usr/share/julia/stdlib/v1.12/REPL/src/REPL.jl:452
#start_repl_backend#41 at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/usr/share/julia/stdlib/v1.12/REPL/src/REPL.jl:427
start_repl_backend at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/usr/share/julia/stdlib/v1.12/REPL/src/REPL.jl:424 [inlined]
#run_repl#50 at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/usr/share/julia/stdlib/v1.12/REPL/src/REPL.jl:653
run_repl at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/usr/share/julia/stdlib/v1.12/REPL/src/REPL.jl:639
jfptr_run_repl_19665.1 at /home/numerik/amartine/.julia/juliaup/julia-1.12.6+0.x64.linux.gnu/share/julia/compiled/v1.12/REPL/u0gqU_E4m7X.so (unknown line)
run_std_repl at ./client.jl:478
jfptr_run_std_repl_24985.1 at /home/numerik/amartine/.julia/juliaup/julia-1.12.6+0.x64.linux.gnu/lib/julia/sys.so (unknown line)
jl_apply at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/julia.h:2391 [inlined]
jl_f_invokelatest at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/builtins.c:881
run_main_repl at ./client.jl:499
repl_main at ./client.jl:586 [inlined]
_start at ./client.jl:561
jfptr__start_63319.1 at /home/numerik/amartine/.julia/juliaup/julia-1.12.6+0.x64.linux.gnu/lib/julia/sys.so (unknown line)
jl_apply at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/julia.h:2391 [inlined]
true_main at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/jlapi.c:971
jl_repl_entrypoint at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/src/jlapi.c:1139
main at /cache/build/builder-amdci5-4/julialang/julia-release-1-dot-12/cli/loader_exe.c:58
unknown function (ip: 0x7f916469a249) at /lib/x86_64-linux-gnu/libc.so.6
__libc_start_main at /lib/x86_64-linux-gnu/libc.so.6 (unknown line)
unknown function (ip: 0x4010b8) at /workspace/srcdir/glibc-2.17/csu/../sysdeps/x86_64/start.S
Allocations: 42089800 (Pool: 42089498; Big: 302); GC: 26
zsh: segmentation fault (core dumped)  julia -g2 --project