Idea to make Zygote support mutation in easy cases

[Lilith]

@mcabbott

What this still won’t solve is that, when the return value of fill!(xs, y) is discarded, its pullback won’t get the gradient for the new xs.

I believe this is an example of what you are referring to, and yes, with my original proposal it would return [0, 0, 0]

gradient([1, 2, 3]) do x
    y = x*x
    fill!(x, zero(x))
    y
end

I don’t fully understand your point (2) above, but for point (1), the standard Julia compiler only elides functions whose return values are not used if they are free of side effects, could Zygote do a similar thing in pullbacks? That is, only elide pullbacks of functions whose return values are ignored when the pullback has no side effects. The trivial pullback of functions with @nograd is free of side effects.

@mcabbott

Now you seem to be saying that the object with this extra flag is present on the forward pass. In which case you can simulate the effect by having the pullback write NaN into the original a. (Or just restore the original values before mutation.)

Yes, this is closer to what I am suggesting. Unfortunately, writing NaN is not a viable option because not all arrays support NaN and restoring the original values is not ideal because that would require all mutating operations to make a copy. I don’t want sum(Float32[x for _ in 1:3]) to allocate twice: first for the array that is populated and summed and then a copy of that uninitialized array to restore later. Ideally there would be some way to poison these arrays at compile time.