Code rewriting can spoil logpdf calculations on GPU

[vargonis]

This issue regards branch rv_gpu. Consider the following:

using Soss
using CuArrays
CuArrays.allowscalar(false)

function gpu(trace)
    d = Dict(k => trace[k] for k in keys(trace))
    for (k,v) in d
        v isa Array ? d[k] = CuArray(v) : nothing
    end
    (; d...)
end

m = @model begin
    # U ~ Uniform()
    X ~ For((i,j)->(1/i,1/j), 2, 3; D=Normal)
end

t = rand(m())
@show logpdf(m(), gpu(t))

This runs OK. However, as soon as the line # U ~ Uniform() is uncommented, one gets an error like:

scalar getindex is disallowed

Stacktrace:
[1] error(::String) at ./error.jl:33
[2] assertscalar(::String) at /home/vargonis/.julia/packages/GPUArrays/1wgPO/src/indexing.jl:14
[3] getindex(::CuArray{Float64,2,Nothing}, ::Int64) at /home/vargonis/.julia/packages/GPUArrays/1wgPO/src/indexing.jl:54
[4] _getindex at ./abstractarray.jl:1004 [inlined]
[5] _broadcast_getindex at ./broadcast.jl:602 [inlined]
[6] getindex at ./broadcast.jl:563 [inlined]
[7] copy(::Base.Broadcast.Broadcasted{Base.Broadcast.ArrayStyle{CuArray},Tuple{Base.OneTo{Int64},Base.OneTo{Int64}},typeof(Soss.cunormlogpdf),Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{2},Nothing,GeneralizedGenerated.Closure{function = (M, i, j;) -> begin
(M.:/(1, i), M.:/(1, j))
end,Tuple{Module}},Tuple{SubArray{Int64,1,UnitRange{Int64},Tuple{Base.Slice{Base.OneTo{Int64}}},true},SubArray{Int64,2,UnitRange{Int64},Tuple{Array{CartesianIndex{0},1},Base.Slice{Base.OneTo{Int64}}},false}}},CuArray{Float64,2,Nothing}}}) at ./broadcast.jl:853
[8] materialize at ./broadcast.jl:819 [inlined]
[9] logpdf(::Product{Normal,GeneralizedGenerated.Closure{function = (M, i, j;) -> begin
(M.:/(1, i), M.:/(1, j))
end,Tuple{Module}},Tuple{SubArray{Int64,1,UnitRange{Int64},Tuple{Base.Slice{Base.OneTo{Int64}}},true},SubArray{Int64,2,UnitRange{Int64},Tuple{Array{CartesianIndex{0},1},Base.Slice{Base.OneTo{Int64}}},false}}}, ::CuArray{Float64,2,Nothing}) at /home/vargonis/.julia/dev/Soss/src/distributions/product.jl:49
[10] _logpdf(::Type{TypeEncoding(Main)}, ::Model{NamedTuple{(),T} where T<:Tuple,TypeEncoding(begin
X ~ For(((i, j)->begin
(1 / i, 1 / j)
end), 2, 3; D=Normal)
U ~ Uniform()
end),TypeEncoding(Main)}, ::NamedTuple{(),Tuple{}}, ::NamedTuple{(:U, :X),Tuple{Float64,CuArray{Float64,2,Nothing}}}) at /home/vargonis/.julia/packages/GeneralizedGenerated/EBwdX/src/closure_conv.jl:54
[11] logpdf(::Soss.JointDistribution{NamedTuple{(),Tuple{}},NamedTuple{(),T} where T<:Tuple,TypeEncoding(begin
X ~ For(((i, j)->begin
(1 / i, 1 / j)
end), 2, 3; D=Normal)
U ~ Uniform()
end),TypeEncoding(Main)}, ::NamedTuple{(:U, :X),Tuple{Float64,CuArray{Float64,2,Nothing}}}) at /home/vargonis/.julia/dev/Soss/src/primitives/logpdf.jl:5
[12] top-level scope at show.jl:562
[13] top-level scope at In[2]:19