Add a flag to control the use of views in gradients

[penelopeysm]

using LogDensityProblems.logdensity_and_gradient with views in e.g.

AdvancedVI.jl/src/algorithms/fisherminbatchmatch.jl

Lines 94 to 97 in 9e6948e

	for b in 1:n_samples
	logπb, gb = LogDensityProblems.logdensity_and_gradient(prob, view(z, :, b))
	grad_buf[:, b] = gb
	logπ_sum += logπb

is incompatible with DynamicPPL.LogDensityFunction because its prepare_gradient is run with a Vector argument. if you then attempt to call logdensity_and_gradient with a SubArray it will error like

  Got exception outside of a @test
  PreparationMismatchError (inconsistent types between preparation and execution):
    - f: ✅
    - backend: ✅
    - x: ❌
      - prep: Vector{Float64}
      - exec: SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}
    - contexts: ✅
  If you are confident that this check is superfluous, you can disable it by running preparation with the keyword argument `strict=Val(false)` inside DifferentiationInterface.

see e.g. https://github.com/TuringLang/Turing.jl/actions/runs/19878777583/job/56972119769?pr=2702.

This PR therefore offers a way to disable this by means of overloading a method on prob. Unfortunately, this is not a very clean solution. it would be nicer to have a field on prob. However, since the only requirement is that prob is an object that satisfies the LogDensityProblems interface, I don't think we can do any better.

It could be put into a keyword argument of algorithm, but I think that's wrong as this is a property of the prob.

[github-actions]

AdvancedVI.jl documentation for PR #223 is available at:
https://TuringLang.github.io/AdvancedVI.jl/previews/PR223/

[penelopeysm]

What I don't really understand is, the same failure should have happened with DynamicPPL 0.38, because gradient prep there was also done with a Vector... So I don't know why CI is failing on TuringLang/Turing.jl#2702 but not TuringLang/Turing.jl#2699 :/

[penelopeysm]

Ahh, it's because [email protected] has an extra line in logdensity_and_gradient which says x = [val for val in x], thus negating the view.

https://github.com/TuringLang/DynamicPPL.jl/blob/6c615ad41946ffa11698408948f85cd4606262ff/src/logdensityfunction.jl#L284-L286

So an alternative fix would be to add that back to DynamicPPL. It doesn't cause issues elsewhere because I think everywhere else that we use LogDensityFunction we just pass it a vector.

An alternative fix would be that, in DynamicPPL we could add a specialisation:

function LogDensityProblems.logdensity_and_gradient(ldf, x::SubArray)
    LogDensityProblems.logdensity_and_gradient(ldf, [v for v in x])
end

function LogDensityProblems.logdensity_and_gradient(ldf, x::AbstractVector{<:Real})
    ...
end

That might actually be easier.

[mhauru]

The specialisation on SubArray looks like a pretty clean solution. Maybe even more general and simpler could be a call to convert, to make any input match the type of prep. This works:

julia> arr = randn(3,3,3);

julia> convert(Vector{Float64}, @view arr[1, :, 1])
3-element Vector{Float64}:
  1.7808290397025959
 -0.5444292197945666
 -0.04197939645098293

It seems like a waste though, to copy the view just because the gradient prep isn't ready for a view. Would it be overkill to make multiple prep objects for different types of inputs? If we really want to engineer this to death, prep could be a mapping of types to prep-objects, that starts out being empty, and every time logdensity_and_gradient is called, if there isn't a prep object for the argument type, one is made and stored for future use. This might be a type stability nightmare though, since we can't change the type of prob in logdensity_and_gradient(prob, x).

EDIT: Or rather, type stability itself would probably be fine, but the cost to pay would be that every logdensity_and_gradient call would result in a runtime look-up of whether a prep object already exists for this argument type. Regardless of whether it does, a prep object of a fixed type would always be returned, it's just a question of whether one was cached already.

[mhauru]

That sounds like a quite a lot of work to try, for uncertain pay-off. It may well be that once the gradient computation starts in earnest the view will have to get copied regardless. So very happy to just copy for now, and see if this can be improved later.

[penelopeysm]

Yup, agreed on all counts.