Zygote Batch Support

[agerlach]

Closes #37

[ChrisRackauckas]

I'd be weary of merging this without figuring out how this is actually different: it looks like it's computing the same thing, just without the parallelism, which seems wrong.

[agerlach]

"wrong" as in there is a better way, or "wrong" as in incorrect? It is computing the same thing. It is equivelent to

dfdp = function (x,p)
    _,back = Zygote.pullback(p->prob.f(x,p),p)
    back(y)[1]
end

if prob.batch > 0
    _dfdp = dfdp
    dfdp = function (x,p)
        out = zeros(length(p),size(x,2))
        for idx in 1:size(x,2)
            out[:,idx] .= _dfdp(@view(x[:,idx]),p)
        end
        out
    end
end

[ChrisRackauckas]

But the first is computing incorrect values right? I think it's best to figure out why, instead of trying to just avoid the issue.

[agerlach]

If you call

dfdp = function (x,p)
    _,back = Zygote.pullback(p->prob.f(x,p),p)
    back(y)[1]
end

directly for batch >0 you will get incorrect answers. I seems like the issue is the size of y. It is odd, in some of the tests it will report the wrong answer, but in others it will error about dimensions. I think the difference between the two depended on if it was inplace or not.

[agerlach]

Here is a MWE showing the issue

using Zygote
w(x,p) = x[1,:]*p[1] + x[2,:]*p[2]
batch = 5
_,back  = Zygote.pullback(p->w(repeat([1,2],inner=(1,batch)),p),[3 2])
back([1])[1] #returns [5, 10]. but we need behaviour like [1 1 1 1 1; 2 2 2 2 2]

[agerlach]

Still feels like a hack, but the pullback can get taken out of the loop w/o the view and then loop through back accordingly.

Using MWE:

using Zygote
w(x,p) = x[1,:]*p[1] + x[2,:]*p[2]
batch = 5
_,back  = Zygote.pullback(p->w(repeat([1,2],inner=(1,batch)),p),[3 2])

out = zeros(2, batch )
for idx in 1:batch
    z = zeros(batch)
    z[idx] = 1
    out[:,idx] = back(z)[1]
end

[ChrisRackauckas]

You shouldn't need to pull back the pieces separately though: this is pullback of the identity matrix of a sparse Jacobian which is the matrix coloring problem that has a single color solution.

[agerlach]

I'm not sure I understand.

[ChrisRackauckas]

it's https://mitmath.github.io/18337/lecture9/stiff_odes but with reverse mode.

[ChrisRackauckas]

using Zygote
w(x,p) = x[1,:]*p[1] + x[2,:]*p[2]
batch = 5
_,back  = Zygote.pullback(p->w(repeat([1,2],inner=(1,batch)),p),[3 2])

out = zeros(2, batch )
z = zeros(batch)
for idx in 1:batch
    z[idx] = 1
    out[:,idx] = back(z)[1]
    z[idx] = 0
end
out

is how I'd do it. It's not a hack but the right way to do it, because you have to pull back the basis elements if the function isn't sparse, and here it's not sparse in p. This keeps the primal calculation down to a single time and is the equivalent to the forwarddiff chunk version.

[agerlach]

This should reduce the primal calculation as we discussed for in and out of place. Currently crashes for nout > 1