vanishing jacobian wrt arg which is a vector of vectors

[lxvm]

Hello,

I noticed an unusual silent error when trying to compute the jacobian of a function with scalar output when the input is an array of SVectors. While gradient and pullback give the expected results, jacobian results an Any array full of zeros.

This is my mwe

using Zygote
using StaticArrays

function mwe(arr)
    tmp = [x[1] for x in arr]
    abs(tmp[1])
end

arr = fill(SVector(1.0, 2.0), 10)
g, = Zygote.gradient(mwe, arr) # OK
j, = Zygote.jacobian(mwe, arr) # wrong
z, zback = Zygote.pullback(mwe, arr) # OK
dz, = zback(1.0)
@show g j dz
# g = SVector{2, Float64}[[1.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]
# j = Any[0 0 0 0 0 0 0 0 0 0]
# dz = SVector{2, Float64}[[1.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]

I used Julia v1.11.6, Zygote v0.7.10, and StaticArrays v1.9.14.

I noticed that the array of Any is happening because type promotion between the scalar output and SVector array element type fails here

Zygote.jl/src/lib/grad.jl

Line 183 in e0af1a8

T = promote_type(eltype(x), eltype(y))

While for my actual calculation gradient is the right function to call, I spent a long time trying figure out if Zygote was failing to differentiate my code due to this silent error. Perhaps it would be worth adding a warning to the docs that only AbstractArray arguments with numeric element type are expected to work?

[mcabbott]

It only understands numbers and arrays of numbers. Essentially because it's not very clear what type it should return in other cases -- what is the jacobian of some struct? What type or shape of object were you hoping for here?

The documentation at present says this:

help?> jacobian
search: jacobian withjacobian factorial join

  jacobian(f, args...) -> Tuple

  For each array a ∈ args this returns a matrix with Ja[k,i] = ∂y[k]/∂a[i] where y = f(args...) is
  usually a vector. Arrays of higher dimension are treated like vec(a), or vec(y) for output.

  For scalar x::Number ∈ args, the result is a vector Jx[k] = ∂y[k]/∂x, while for scalar y all
  results have just one row.

  With any other argument type, no result is produced, even if gradient would work.

[...]

  │ Warning
  │
  │  For arguments of any type except Number & AbstractArray, the result is nothing.

That's not very specific that "array" and "AbstractArray" mean AbstractArray{<:Number}, sadly.

The fact that you get a zero Matrix{Any} instead of nothing probably indicates that the code is also insufficiently careful about this distinction.

[lxvm]

I would support narrowing the definition of jacobian to have arguments of AbstractArray{<:Number}s and to throw an error in other cases. My expectation, however, was just that the structural tangents would be returned, similar to gradient and pullback, since I was using an array of SVectors as the argument to a function that evaluates a multivariate polynomial for each array element that corresponds to a quadrature node in a convolution integral that is part of my loss function. I am frequently using arrays of StaticArrays in this way to represent vector fields on discretized geometries and have been trying to use the same data type throughout the code, although I realize that this type can also be reinterpreted as an array with an additional dimension containing the elements of the StaticArray.

[mcabbott]

The function mwe returns a scalar, so you can use gradient, and structural tangents are simple. The question about the return type of jacobian is really for functions which return an array. It has to add one more dimension, and the question is where.

That is, jacobian(nt -> nt.x, (x = rand(3), y = rand(3))) could return either something like (x=matrix(3,3), y=nothing), or something like [(x=vector(3), y=nothing)) for _ in 1:3]. The code for the first case is much easier to write, but also seems unlike the jacobian(cumsum, rand(3)) case, which does not return a vector of vectors. With arrays of arrays, jacobian(first, [rand(3) for _ in 1:2]) has a similar problem.

[lxvm]

If I understand correctly, if f is a function that maps an input array of length N to an output array of length M then jacobian will return an M by N matrix of the partial derivatives and will always flatten the input and output arrays to ignore their shape. If the input array is an array of structs, then each element of the Jacobian can be a structural tangent because that represents a gradient of a matrix element of f w.r.t. the struct. This should match the latter of the return values you suggested in your example, which is most similar to how the code is currently written. However, the current code fails to infer the type of the structural tangent because of the promote_type call cited in the O.P. and the only fix I can think of is to change the code to not preallocate the output arrays of a fixed type and instead dynamically fill the Jacobian row by row with the structural tangents returned by pullback. I can give this a shot although I wonder how not being restrictive about types will affect type stability if the type of the output array needs to be widened.

[lxvm]

Wait, sorry, clearly the promote_type call returns Any which is a wide enough type for the structural tangent. I think I can narrow down the error to _gradcopy! at the following method

Zygote.jl/src/lib/grad.jl

Line 214 in e0af1a8

_gradcopy!(dst::AbstractArray, src::AbstractArray) = copyto!(dst, g isa Number ? g : 0 for g in src) # e.g. Union{Nothing,Float64}

This line will fill the output array with zeros when the input array elements are anything but numbers. However, I think this line intends to put zeros only where there are nothings and it could still work if there were a way to fill it with zero(eltype(dest)). I would change the ternary statement to something like g === nothing ? zero(eltype(dest)) : g although I don't know how breaking a change that could be. It would also clearly fail for mwe because the element type of the output would be Any which doesn't have a zero method

[mcabbott]

There are many implementation details!

For now #1582 fixes the bug that it does not, at present, return nothing here.

I'm not really sold that we should extend jacobian beyond arrays of numbers. It has one job, and can have a clear description. If you want to create arrays of structural tangents, the way you can do it now is like so:

julia> arg = (x = rand(3), y = rand(3));

julia> out, back = Zygote.pullback(nt -> nt.x, arg);

julia> back.(eachcol(I(3)))
3-element Vector{Tuple{@NamedTuple{x::Vector{Float64}, y::Nothing}}}:
 ((x = [1.0, 0.0, 0.0], y = nothing),)
 ((x = [0.0, 1.0, 0.0], y = nothing),)
 ((x = [0.0, 0.0, 1.0], y = nothing),)

That's not one line. But it's also probably clearer than reading up on special cases of jacobian that don't return matrices of size output-by-input?

[lxvm]

I have no objection to restricting jacobian to (multi-argument) functions from R^N1 x R^N2 x ... to R^M as that is the standard setting for which it is defined. I also think #1582 is useful because it makes these assumptions more explicit. In particular for the mwe, getting nothing for a Jacobian is an improvement over getting arrays of zeros because for debugging the former is more likely to make someone think the function isn't the right one to call whereas the latter leads one to think there is some kind of silent bug.

Thank you as well for the simple example for accumulating Jacobians w.r.t. arrays of structs. It captures the essence of what jacobian does minus the handling of multiple arguments and reshaping+concatenation of gradients

Fixed by #1582