Gradients dropped by adapt

[mcabbott]

Moving y to "gpu" inside loss causes its gradient to be lost:

julia> using Tracker, JLArrays

julia> JLArrays.allowscalar(false)

julia> Tracker.withgradient((x,y) -> sum(x[1:2] + jl(y))^2, jl([1,2,3.0]), [4,5.0])
(val = 144.0, grad = ([24.0, 24.0, 0.0], [0.0, 0.0])) 

julia> ans.grad[1] isa JLArray
true

unlike Zygote:

julia> Zygote.withgradient((x,y) -> sum(x[1:2] + jl(y))^2, jl([1,2,3.0]), [4,5.0])
(val = 144.0, grad = ([24.0, 24.0, 0.0], [24.0, 24.0]))

[ToucheSir]

Riffing on a rule:

@grad function Adapt.adapt_storage(adaptor, x::AT) where {AT <: Array}
  adapt_storage_pullback(Δ) = (nothing, Adapt.adapt_storage(AT, Δ))
  return Adapt.adapt_storage(adaptor, x), adapt_storage_pullback
end