Zygote is broken for Stacked transformations

[bgctw]

The following MWE fails for Bijectors Bijectors v0.15.3

#]activate --temp
#add Bijectors, Zygote
using Bijectors
using Zygote
x = [-2, 3.0, 2.1, 2.2]
ranges = [1:1, 2:2, 3:4]
tr = Stacked((elementwise(identity), elementwise(exp), elementwise(exp)), ranges)
y = tr(x)
gr = Zygote.gradient(tr, x) # error: output length mismatch (4 != 3)

Environment

julia> Pkg.status()
Status `/tmp/jl_fDy5Lg/Project.toml`
  [76274a88] Bijectors v0.15.3
  [e88e6eb3] Zygote v0.7.1

[mhauru]

This should be fixed by #362 once it's merged. Thanks for spotting and reporting @bgctw!

Note that since tr in your example returns a vector, you need to call Zygote.jacobian rather than Zygote.gradient. With #362 this works:

julia> using Bijectors
       using Zygote
       x = [-2, 3.0, 2.1, 2.2]
       ranges = [1:1, 2:2, 3:4]
       tr = Stacked((elementwise(identity), elementwise(exp), elementwise(exp)), ranges)
       y = tr(x)
       Zygote.jacobian(tr, x)
([1.0 0.0 0.0 0.0; 0.0 20.085536923187668 0.0 0.0; 0.0 0.0 8.16616991256765 0.0; 0.0 0.0 0.0 9.025013499434122],)

[mhauru]

Bijectors v0.15.4 should be out later today, and includes the fix for this.

[bgctw]

Thanks for this quick fix and the quick release of it.