Jacobians for functions beyond vector-to-vector?

[prbzrg]

julia> using AbstractDifferentiation, Zygote

julia> function ft1(x)
       y1 = (x[1]*x[2])^x[3]
       y2 = (x[2]*x[3])^x[1]
       y3 = (x[3]*x[1])^x[2]
       [y1, y2, y3]
       end
ft1 (generic function with 1 method)

julia> function ft2(xs)
       ys = ft1.(eachcol(xs))
       hcat(ys...)
       end
ft2 (generic function with 1 method)

julia> r3 = rand(3, 8)
3×8 Matrix{Float64}:
 0.0354617  0.444021  0.161892  0.56656   0.92774   0.260982  0.839223  0.175217
 0.020074   0.185554  0.747159  0.850257  0.930541  0.451429  0.978923  0.937234
 0.213358   0.838412  0.562181  0.256845  0.743921  0.777094  0.207115  0.791544

julia> only(Zygote.jacobian(ft2, r3))
24×24 Matrix{Float64}:
  1.28169   2.26417  -1.54394     0.0        0.0        0.0        0.0        0.0       …   0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
 -4.4943    1.45594   0.136984    0.0        0.0        0.0        0.0        0.0           0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
  0.51321  -4.42796   0.0852995   0.0        0.0        0.0        0.0        0.0           0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.232868   0.557241  -0.307858   0.0        0.0           0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0        -0.814452   1.04745    0.231817   0.0        0.0           0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.347887  -0.822595   0.18424    0.0        0.0       …   0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        1.05908    0.229477      0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0       -0.753767   0.188289      0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        0.769966  -0.39986       0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0           0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0       …   0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0           0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0           0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0           0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0           0.0        0.0        0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0       …   0.326677  -0.405857   0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0           0.439836   0.255509   0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0          -0.776414   0.282692   0.0        0.0        0.0        0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0           0.0        0.0        0.236948   0.203133  -0.188738   0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0           0.0        0.0       -0.418176   0.224654   1.06182    0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0       …   0.0        0.0        0.210367  -0.315561   0.852398   0.0        0.0        0.0
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0           0.0        0.0        0.0        0.0        0.0        1.08112    0.202117  -0.432339
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0           0.0        0.0        0.0        0.0        0.0       -0.283372   0.177422   0.210078
  0.0       0.0       0.0         0.0        0.0        0.0        0.0        0.0           0.0        0.0        0.0        0.0        0.0        0.839794  -0.310155   0.185898

julia> only(AD.jacobian(AD.ZygoteBackend(), ft2, r3))
ERROR: "The function `identity_matrix_like` is not defined for the type Matrix{Float64}."
Stacktrace:
 [1] identity_matrix_like(x::Matrix{Float64})
   @ AbstractDifferentiation C:\Users\Hossein Pourbozorg\.julia\packages\AbstractDifferentiation\o62DE\src\AbstractDifferentiation.jl:612
 [2] jacobian(ab::AbstractDifferentiation.ReverseRuleConfigBackend{Zygote.ZygoteRuleConfig{Zygote.Context}}, f::Function, xs::Matrix{Float64})
   @ AbstractDifferentiation C:\Users\Hossein Pourbozorg\.julia\packages\AbstractDifferentiation\o62DE\src\AbstractDifferentiation.jl:570
 [3] top-level scope
   @ REPL[22]:1
 [4] top-level scope
   @ C:\Users\Hossein Pourbozorg\.julia\packages\CUDA\tTK8Y\src\initialization.jl:52

[gdalle]

I think that's due to the Jacobian being an ill-defined notion when your function is not vector-to-vector. Many backends define it by flattening, we probably should do that in AD too

[olivierlabayle]

Hi, I am having a similar issue with the vcat function, it may also be related to a few issues I've seen raised.

I think I'm following the docs here but:

import AbstractDifferentiation as AD, Zygote

backend = AD.ZygoteBackend()
args = (1, 2)

# This works fine
jac_from_ad = AD.jacobian(backend, (x, y) -> vcat(x,y), args...)
jac_from_zyg = Zygote.jacobian(vcat, args...)
@assert jac_from_ad == jac_from_zyg

# But this does not evaluate
failed = AD.jacobian(backend, vcat, args...)