fix JET infer

[thautwarm]

fix JuliaStaging/GeneralizedGenerated.jl#69 : Fix the following tests:

using Soss, JET

m1 = @model N begin
    p ~ Uniform()
    x ~ For(N) do j
            Bernoulli(p / j)
        end
    end

@test_opt rand(m1(10))

m3 = @model N begin
    p ~ Uniform()
    f(ctx) = Base.Fix1(ctx) do ctx, j
        Bernoulli(ctx.p / j)
    end
    x ~ For(f((p=p,)), N)
end

@test_opt rand(m3(10))