Gradients of functions that use Buffer don't pass the inferred test

[gabrevaya]

I don't know if this behavior is intended or maybe related to #1355, but just in case I report it here.

The return type of the function vstack from the following MWE can be successfully inferred but not the one from its gradient:

using Zygote
using Test

function vstack(xs)
    buf = Zygote.Buffer(xs, length(xs), 5)
    for i = 1:5
      buf[:, i] = xs
    end
    return copy(buf)
end

@inferred vstack([1, 2, 3]);
@inferred gradient(x -> sum(vstack(x)), [1, 2, 3]);

ERROR: return type Tuple{Vector{Float64}} does not match inferred return type Tuple{Any}
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] top-level scope
   @ ~/Documents/doctorado/issues/Zygote_buffer_gradient/issue.jl:13

(Zygote_buffer_gradient) pkg> st
Status `~/Documents/doctorado/issues/Zygote_buffer_gradient/Project.toml`
  [e88e6eb3] Zygote v0.6.59

julia> versioninfo()
Julia Version 1.8.5
Commit 17cfb8e65ea (2023-01-08 06:45 UTC)
Platform Info:
  OS: macOS (arm64-apple-darwin21.5.0)
  CPU: 8 × Apple M1 Pro
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-13.0.1 (ORCJIT, apple-m1)
  Threads: 6 on 6 virtual cores
Environment:
  JULIA_EDITOR = code

[ToucheSir]

It's neither intended nor due to a previous PR, but a limitation of Zygote internals. Like other mutable values (e.g. mutable structs and Dicts), manipulating Buffers will lead to type unstable gradient code. That's part of why we recommend avoiding them wherever possible.

[gabrevaya]

Ok, thanks!