Non-determinstic segfault

[ptiede]

I don't have a true MWE at this point and it only segfaults sporadically but here is a MWE

using Comrade
using Distributions
using VLBIImagePriors
using Enzyme
using LogDensityProblemsAD
using Serialization
using TypedTables
# ## Load the Data

logistic(x) = inv(1.0 + exp(-x))
function sky(θ, metadata)
    (;c, σ, p, p0, pσ, angparams) = θ
    (;ftot, grid) = metadata
    ## Build the stokes I model
    rast = to_simplex(CenteredLR(), c.params*σ)
    rast .= ftot.*rast
    ## The total polarization fraction is modeled in logit space so we transform it back
    pim = logistic.(p0 .+ pσ.*p.params)
    ## Build our IntensityMap
    pmap = PoincareSphere2Map(rast, pim, angparams, grid)
    ## Construct the actual image model which uses a third order B-spline pulse
    m = ContinuousImage(pmap, BSplinePulse{3}())
    ## Finally find the image centroid and shift it to be at the center
    x0, y0 = centroid(pmap)
    ms = shifted(m, -x0, -y0)
    return ms
end

dvis = deserialize(joinpath(@__DIR__, "data.jls"))
        

fovx = μas2rad(60.0)
fovy = μas2rad(60.0)
nx = ny = 64
grid = imagepixels(fovx, fovy, nx, ny)
skymeta = (;ftot=1.0, grid)
rat = beamsize(dvis)/step(grid.X)
cmarkov = ConditionalMarkov(GMRF, grid; order=1)
dρ = truncated(InverseGamma(1.0, -log(0.1)*rat); lower=2.0, upper=max(nx, ny))
cprior = HierarchicalPrior(cmarkov, dρ)

fwhmfac = 2.0*sqrt(2.0*log(2.0))
skyprior = (
    c = cprior,
    σ  = truncated(Normal(0.0, 0.1); lower=0.0),
    p  = cprior,
    p0 = Normal(-2.0, 2.0),
    pσ =  truncated(Normal(0.0, 1.0); lower=0.01),
    angparams = ImageSphericalUniform(nx, ny),
    )

skym = SkyModel(sky, skyprior, grid; metadata=skymeta)

function fgain(x)
    gR = exp(x.lgμ +  x.lgσ*x.lgR + 1im*x.gpR)
    lgrat = x.lgratμ + x.lgratσ*x.lgrat
    gprat = x.gpratμ + x.gpratσ*x.gprat
    gL = gR*exp(lgrat + 1im*gprat)
    return gR, gL
end
G = JonesG(fgain)
function fdterms(x)
    dR = complex(x.dRx, x.dRy)
    dL = complex(x.dLx, x.dLy)
    return dR, dL
end

D = JonesD(fdterms)

R = JonesR(;add_fr=true)
J = JonesSandwich(*, G, D, R)


intprior = (
    lgR    = ArrayPrior(IIDSitePrior(ScanSeg(), Normal(0.0, 1.0))),
    lgμ    = ArrayPrior(IIDSitePrior(TrackSeg(), Normal(0.0, 0.1))),
    lgσ    = ArrayPrior(IIDSitePrior(TrackSeg(), Exponential(0.1))),
    gpR    = ArrayPrior(IIDSitePrior(ScanSeg(), DiagonalVonMises(0.0, inv(π  ^2))); refant=SEFDReference(0.0), phase=true),
    lgrat  = ArrayPrior(IIDSitePrior(ScanSeg(), Normal(0.0, 1.0))),
    lgratμ = ArrayPrior(IIDSitePrior(TrackSeg(), Normal(0.0, 0.1))),
    lgratσ = ArrayPrior(IIDSitePrior(TrackSeg(), Exponential(0.1))),
    gprat  = ArrayPrior(IIDSitePrior(ScanSeg(), DiagonalVonMises(0.0, 1.0)); refant = SingleReference(:AA, 0.0)),
    gpratμ = ArrayPrior(IIDSitePrior(TrackSeg(), DiagonalVonMises(0.0, 0.1)); refant = SingleReference(:AA, 0.0)),
    gpratσ = ArrayPrior(IIDSitePrior(TrackSeg(), Exponential(0.1)); refant = SingleReference(:AA, 0.0)),
    dRx    = ArrayPrior(IIDSitePrior(TrackSeg(), Normal(0.0, 0.2))),
    dRy    = ArrayPrior(IIDSitePrior(TrackSeg(), Normal(0.0, 0.2))),
    dLx    = ArrayPrior(IIDSitePrior(TrackSeg(), Normal(0.0, 0.2))),
    dLy    = ArrayPrior(IIDSitePrior(TrackSeg(), Normal(0.0, 0.2))),
)
intmodel = InstrumentModel(J, intprior)


post = VLBIPosterior(skym, intmodel, dvis)
tpost = asflat(post)

x = prior_sample(tpost)
dx = zero(x)
autodiff(Reverse, logdensityof, Active, Const(tpost), Duplicated(x, dx))

function trytosegfault(tpost, x, dx, n=1_000_000_000)
    ℓ = ADgradient(Val(:Enzyme), tpost)
    for _ in 1:n
        LogDensityProblemsAD.logdensity_and_gradient(ℓ, x)
    end
    return dx
end
trytosegfault(tpost, x, dx)

This only segfaults sometimes. I am still trying to reduce the code and make it more consistent but it hasn't been simple.
Note you need to be on the Comrade#ptiede-enzymegc branch.

[wsmoses]

@ptiede can you check if current main hits this still?