Relax Enzyme bounds

[ptiede]

No description provided.

[ptiede]

Enzyme sometimes fails and sometimes passes. Who knows why. This only seems to happen during tests so I am going to play with things until I figure out why.

[ptiede]

So locally, the tests fail non-deterministically. So that is great...

[ptiede]

Ok this consistently fails. The trick is I needed to evaluate the gradient twice at the same location

function test_nonanalytic(intm, algorithm, vis)
    skym = SkyModel(test_model, test_prior(), imagepixels(μas2rad(150.0), μas2rad(150.0), 256, 256); algorithm)
    post = VLBIPosterior(skym, intm, vis)

    tpostf = asflat(post)
    x0 = prior_sample(tpostf)
    @inferred logdensityof(tpostf, x0)
    @info logdensityof(tpostf, x0)
    f0, gz = LogDensityProblems.logdensity_and_gradient(tpostf, x0)
    # @info f0
    f0, gz = LogDensityProblems.logdensity_and_gradient(tpostf, x0)
    @info f0
    mfd = central_fdm(5, 1)
    gn, = FiniteDifferences.grad(mfd, tpostf, x0)
    return @test gz ≈ gn
end



function test_model(θ, meta)
    m1 = θ.f1 * rotated(stretched(Gaussian(), θ.σ1 * θ.τ1, θ.σ1), θ.ξ1)
    m2 = θ.f2 * rotated(stretched(Gaussian(), θ.σ2 * θ.τ2, θ.σ2), θ.ξ2)
    return m1 + shifted(m2, θ.x, θ.y)
end

function test_prior()
    return (
        f1 = Uniform(0.8, 1.2),
        σ1 = Uniform(μas2rad(1.0), μas2rad(40.0)),
        τ1 = Uniform(0.35, 0.65),
        ξ1 = Uniform(-π / 2, π / 2),
        f2 = Uniform(0.3, 0.7),
        σ2 = Uniform(μas2rad(1.0), μas2rad(40.0)),
        τ2 = Uniform(0.35, 0.65),
        ξ2 = Uniform(-π / 2, π / 2),
        x = Uniform(-μas2rad(40.0), μas2rad(40.0)),
        y = Uniform(-μas2rad(40.0), μas2rad(40.0)),
    )
end

    _, vis, amp, lcamp, cphase = load_data()

fgain(x) = exp.(x.gp)

G = SingleStokesGain(fgain)
intm = InstrumentModel(
        G, (
            gp = ArrayPrior(IIDSitePrior(TrackSeg(), Normal(0.0, 1.0)), refant = SingleReference(:AA, 0.0), phase = false),
        )
    )

@testset "NFFT" begin
   test_nonanalytic(intm, NFFTAlg(), vis[1:2])
end

[ptiede]

Better MWE of gradient issue

using Comrade, Distributions, FiniteDifferences, Enzyme, Test

vis = deserialize("test_data.jls") #ask me for this
G = SingleStokesGain(fgain)
intm = InstrumentModel(G, (;lg = ArrayPrior(IIDSitePrior(TrackSeg(), Normal())),))

# Model is literally a constant
function test_model(θ, meta)
    m1 = ZeroModel()
    return m1
end

# Only piece here that is strange is TransformVariables.
function test_prior()
    return (
        f1 = Uniform(0.8, 1.2),
        σ1 = Uniform(μas2rad(1.0), μas2rad(40.0)),
        τ1 = Uniform(0.35, 0.65),
        ξ1 = Uniform(-π / 2, π / 2),
        f2 = Uniform(0.3, 0.7),
        σ2 = Uniform(μas2rad(1.0), μas2rad(40.0)),
        τ2 = Uniform(0.35, 0.65),
        ξ2 = Uniform(-π / 2, π / 2),
        x = Uniform(-μas2rad(40.0), μas2rad(40.0)),
        y = Uniform(-μas2rad(40.0), μas2rad(40.0)),
    )
end

skym = SkyModel(test_model, test_prior(), imagepixels(μas2rad(150.0), μas2rad(150.0), 256, 256))
post = VLBIPosterior(skym, intm, vis)
tpostf = asflat(post)
x0 = rand(dimension(tpostf))


using TransformVariables

function foo(dp, tr, x)
           y, lj = transform_and_logjac(tr, x)
           return logdensityof(dp, y) + lj
end

trf = tpostf.transform
dp = tpostf.lpost.prior

Enzyme.gradient(set_runtime_activity(Enzyme.Reverse), foo, Const(dp), Const(trf), x0)[3][3]
# 0.9321000954069457

Enzyme.gradient(set_runtime_activity(Enzyme.Reverse), foo, Const(dp), Const(trf), x0)[3][3]
# -0.9321000954069457

Enzyme.gradient(set_runtime_activity(Enzyme.Reverse), foo, Const(dp), Const(trf), x0)[3][3]
# 2.7963002862208373

Enzyme.gradient(set_runtime_activity(Enzyme.Reverse), foo, Const(dp), Const(trf), x0)[3][3]
# -4.660500477034729

Enzyme.gradient(set_runtime_activity(Enzyme.Reverse), foo, Const(dp), Const(trf), x0)[3][3]
# 10.253101049476403