Pullback over jacobian

[aksuhton]

I get a mutation error when I try to do the following:

I have a parameterized function that takes a set of inputs. I would like to take the jacobian of this function (really just the diagonal) with respect to its inputs and then perform gradients on the resulting object with respect to the parameters.

Not long ago I posted about this in a semi-related issue, but my example code at the time involved neural networks. I'm posting here now with an MWE that I think does indeed belong in a Zygote issue.

using Zygote
# Setup
x = randn(Float32, 10, 2)
ps = randn(Float32, 10, 10)
f(x, ps) = exp.(ps * x)
# Forward
function get_jac(f, x, ps)
    jac = only(Zygote.jacobian(x -> f(x, ps), x))
    return jac
end
jac = get_jac(f, x, ps)
# Backward
function get_grads(f, x, ps)
    l, back = Zygote.pullback(p -> get_jac(f, x, p), ps)
    # xxx
    gs = only(back(one.(l)))
    return gs
end
# Mutation error
gs = get_grads(f, x, ps)

The stacktrace

1-element ExceptionStack:
LoadError: Mutating arrays is not supported -- called copyto!(SubArray{Float32, 1, Matrix{Float32}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}, ...)
This error occurs when you ask Zygote to differentiate operations that change
the elements of arrays in place (e.g. setting values with x .= ...)

Possible fixes:
- avoid mutating operations (preferred)
- or read the documentation and solutions for this error
  https://fluxml.ai/Zygote.jl/latest/limitations

Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] _throw_mutation_error(f::Function, args::SubArray{Float32, 1, Matrix{Float32}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/lib/array.jl:70
  [3] (::Zygote.var"#543#544"{SubArray{Float32, 1, Matrix{Float32}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}})(::Nothing)
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/lib/array.jl:85
  [4] (::Zygote.var"#2633#back#545"{Zygote.var"#543#544"{SubArray{Float32, 1, Matrix{Float32}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}}})(Δ::Nothing)
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
  [5] _gradcopy!
    @ ~/.julia/packages/Zygote/jxHJc/src/lib/grad.jl:171 [inlined]
  [6] (::Zygote.Pullback{Tuple{typeof(Zygote._gradcopy!), SubArray{Float32, 1, Matrix{Float32}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}, Matrix{Float32}}, Tuple{Zygote.                     var"#2633#back#545"{Zygote.var"#543#544"{SubArray{Float32, 1, Matrix{Float32}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}}}}})(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface2.jl:0
  [7] withjacobian
    @ ~/.julia/packages/Zygote/jxHJc/src/lib/grad.jl:153 [inlined]
  [8] (::Zygote.Pullback{Tuple{typeof(withjacobian), var"#9#10"{typeof(f), Matrix{Float32}}, Matrix{Float32}}, Any})(Δ::@NamedTuple{val::Nothing, grad::Tuple{Matrix{Float32}}})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface2.jl:0
  [9] #291
    @ ~/.julia/packages/Zygote/jxHJc/src/lib/lib.jl:206 [inlined]
 [10] #2169#back
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72 [inlined]
 [11] jacobian
    @ ~/.julia/packages/Zygote/jxHJc/src/lib/grad.jl:128 [inlined]
 [12] (::Zygote.Pullback{Tuple{typeof(jacobian), var"#9#10"{typeof(f), Matrix{Float32}}, Matrix{Float32}}, Tuple{Zygote.var"#2169#back#293"{Zygote.var"#291#292"{Tuple{Tuple{Nothing}, Tuple{Nothing}}, Zygote.Pullback{Tuple{typeof(withjacobian), var"#9#10"{typeof(f), Matrix{Float32}}, Matrix{Float32}}, Any}}}, Zygote.var"#2180#back#303"{Zygote.var"#back#302"{:grad, Zygote.Context{false}, @          NamedTuple{val::Vector{Float32}, grad::Tuple{Matrix{Float32}}}, Tuple{Matrix{Float32}}}}, Zygote.var"#2013#back#204"{typeof(identity)}}})(Δ::Tuple{Matrix{Float32}})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface2.jl:0
 [13] get_jac
    @ ~/GitHub/DTWINS/scripts/thermo/zygdev.jl:10 [inlined]
 [14] (::Zygote.Pullback{Tuple{typeof(get_jac), typeof(f), Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(only), Tuple{Matrix{Float32}}}, Tuple{Zygote.                          var"#2029#back#213"{Zygote.var"#back#211"{1, 1, Zygote.Context{false}, Matrix{Float32}}}}}, Zygote.Pullback{Tuple{typeof(jacobian), var"#9#10"{typeof(f), Matrix{Float32}}, Matrix{Float32}},           Tuple{Zygote.var"#2169#back#293"{Zygote.var"#291#292"{Tuple{Tuple{Nothing}, Tuple{Nothing}}, Zygote.Pullback{Tuple{typeof(withjacobian), var"#9#10"{typeof(f), Matrix{Float32}}, Matrix{Float32}},      Any}}}, Zygote.var"#2180#back#303"{Zygote.var"#back#302"{:grad, Zygote.Context{false}, @NamedTuple{val::Vector{Float32}, grad::Tuple{Matrix{Float32}}}, Tuple{Matrix{Float32}}}}, Zygote.               var"#2013#back#204"{typeof(identity)}}}, Zygote.var"#2210#back#313"{Zygote.Jnew{var"#9#10"{typeof(f), Matrix{Float32}}, Nothing, false}}}})(Δ::Matrix{Float32})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface2.jl:0
 [15] #11
    @ ~/GitHub/DTWINS/scripts/thermo/zygdev.jl:18 [inlined]
 [16] (::Zygote.Pullback{Tuple{var"#11#12"{typeof(f), Matrix{Float32}}, Matrix{Float32}}, Tuple{Zygote.var"#2180#back#303"{Zygote.var"#back#302"{:f, Zygote.Context{false}, var"#11#12"{typeof(f),      Matrix{Float32}}, typeof(f)}}, Zygote.var"#2180#back#303"{Zygote.var"#back#302"{:x, Zygote.Context{false}, var"#11#12"{typeof(f), Matrix{Float32}}, Matrix{Float32}}}, Zygote.                          Pullback{Tuple{typeof(get_jac), typeof(f), Matrix{Float32}, Matrix{Float32}}, Tuple{Zygote.Pullback{Tuple{typeof(only), Tuple{Matrix{Float32}}}, Tuple{Zygote.var"#2029#back#213"{Zygote.               var"#back#211"{1, 1, Zygote.Context{false}, Matrix{Float32}}}}}, Zygote.Pullback{Tuple{typeof(jacobian), var"#9#10"{typeof(f), Matrix{Float32}}, Matrix{Float32}}, Tuple{Zygote.                        var"#2169#back#293"{Zygote.var"#291#292"{Tuple{Tuple{Nothing}, Tuple{Nothing}}, Zygote.Pullback{Tuple{typeof(withjacobian), var"#9#10"{typeof(f), Matrix{Float32}}, Matrix{Float32}}, Any}}}, Zygote.   var"#2180#back#303"{Zygote.var"#back#302"{:grad, Zygote.Context{false}, @NamedTuple{val::Vector{Float32}, grad::Tuple{Matrix{Float32}}}, Tuple{Matrix{Float32}}}}, Zygote.                              var"#2013#back#204"{typeof(identity)}}}, Zygote.var"#2210#back#313"{Zygote.Jnew{var"#9#10"{typeof(f), Matrix{Float32}}, Nothing, false}}}}}})(Δ::Matrix{Float32})
    @ Zygote ~/.julia/packages/Zygote/jxHJc/src/compiler/interface2.jl:0

Version info

Julia Version 1.10.2                                                                               
Commit bd47eca2c8a (2024-03-01 10:14 UTC)                                                          
Build Info:                                                                                        
  Official https://julialang.org/ release                                                          
Platform Info:                                                                                     
  OS: Linux (x86_64-linux-gnu)                                                                     
  CPU: 32 × AMD Ryzen 9 7950X3D 16-Core Processor                                                  
  WORD_SIZE: 64                                                                                    
  LIBM: libopenlibm                                                                                
  LLVM: libLLVM-15.0.7 (ORCJIT, znver3)                                                            
Threads: 1 default, 0 interactive, 1 GC (on 32 virtual cores)                                      
Environment:                                                                                       
  LD_PRELOAD = /usr/${LIB}/libgtk3-nocsd.so.0  

[ToucheSir]

This is either a duplicate of #1268 or #953.

[aksuhton]

@ToucheSir

Thank you for taking the time to close this issue with links. In #1268, @mcabbott writes

"Zygote's jacobian function isn't Zygote-differentiable. There's no major barrier to making it so, someone just has to do it. "

Would you say that this is still an accurate statement?

[mcabbott]

Yes, it is. Would not be so hard to do really.

[flo-he]

Are there any updates on this? The issue of non-differentiable Jacobians exists now for almost 2 years and given that "there's no major barrier to making it so, someone just has to do it", I wonder why this isn't fixed by now? I think this problem makes it very hard to favor Julia over Jax/Pytorch for more "intricate" deep learning projects.

[avik-pal]

@flo-he solving it in the general case is non-trivial (other than making the jacobian function non-mutating, it requires reverse over reverse to be efficient). You can work around it using https://lux.csail.mit.edu/stable/manual/nested_autodiff which makes certain assumptions on how Zygote.jacobian is being called

[flo-he]

I see, thank you @avik-pal!