Part 1 of autograd backend: DifferentiationInterface extention

.github/workflows/CI.yml

@@ -33,7 +33,7 @@ jobs:
     steps:
       - uses: actions/checkout@v6
 
-      - uses: julia-actions/setup-julia@v2
+      - uses: julia-actions/setup-julia@v3
         with:
           version: ${{ matrix.runner.version }}
 
@@ -47,28 +47,36 @@ jobs:
 
       - uses: julia-actions/julia-processcoverage@v1
 
-      - uses: codecov/codecov-action@v5
+      - uses: codecov/codecov-action@v6
         with:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
           fail_ci_if_error: true
 
   ext:
-    name: Ext (logdensityproblems, ${{ matrix.version }})
+    name: Ext (${{ matrix.label }}, ${{ matrix.version }})
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
+        label:
+          - ext/differentiationinterface
         version:
           - '1'
           - 'min'
     steps:
       - uses: actions/checkout@v6
-      - uses: julia-actions/setup-julia@v2
+      - uses: julia-actions/setup-julia@v3
         with:
           version: ${{ matrix.version }}
       - uses: julia-actions/cache@v3
       - uses: julia-actions/julia-buildpkg@v1
-      - run: julia --project=. test/run_extras.jl
+      - run: julia --code-coverage=user test/run_extras.jl
         env:
-          LABEL: ext/logdensityproblems
+          LABEL: ${{ matrix.label }}
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v6
+        with:
+          files: lcov.info
+          token: ${{ secrets.CODECOV_TOKEN }}
+          fail_ci_if_error: true

AGENTS.md

@@ -24,14 +24,12 @@ AbstractPPL.jl provides shared interfaces and utilities for probabilistic progra
   - When touching `VarName` or optics code, test symbolic, indexed, nested, and serialization round-trip cases.
   - Respect evaluator contracts: `VectorEvaluator` is for flat vectors; `NamedTupleEvaluator` is for stable named structures; `!!` derivative APIs may return cache-aliased arrays.
   - Use `check_dims=false` only for trusted AD hot paths. Public evaluator calls should validate user input.
-  - Keep `LogDensityProblems` integration vector-only unless its contract changes.
 
 ## Tests
 
   - Core tests: `GROUP=Tests julia --project=test test/runtests.jl`
   - Doctests: `GROUP=Doctests julia --project=test test/runtests.jl`
   - Full package tests: `julia --project=. -e 'using Pkg; Pkg.test()'`
-  - LogDensityProblems extension: `LABEL=ext/logdensityproblems julia --project=. test/run_extras.jl`
   - Docs: `julia --project=docs docs/make.jl`
 
 Run the smallest relevant test first, then broaden when changing public interfaces, extensions, or downstream-facing behaviour. Do not weaken tests just to make CI pass.

JULIA.md

@@ -1,6 +1,6 @@
 # JULIA.md
 
-Shared day-to-day Julia practices. DynamicPPL-specific review notes live in `AGENTS.md`; newcomer context lives in `docs/src/onboarding.md`.
+Shared day-to-day Julia practices. AbstractPPL-specific review notes live in `AGENTS.md`.
 
 ## Engineering
 

Project.toml

@@ -1,6 +1,6 @@
 name = "AbstractPPL"
 uuid = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
-keywords = ["probablistic programming"]
+keywords = ["probabilistic programming"]
 license = "MIT"
 desc = "Common interfaces for probabilistic programming"
 version = "0.14.3"
@@ -19,25 +19,28 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [weakdeps]
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
-LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [extensions]
+AbstractPPLDifferentiationInterfaceExt = ["DifferentiationInterface"]
 AbstractPPLDistributionsExt = ["Distributions", "LinearAlgebra"]
-AbstractPPLLogDensityProblemsExt = ["LogDensityProblems"]
+AbstractPPLTestExt = ["Test"]
 
 [compat]
 ADTypes = "1"
-LogDensityProblems = "2"
 AbstractMCMC = "2, 3, 4, 5"
 Accessors = "0.1"
 BangBang = "0.4"
 DensityInterface = "0.4"
+DifferentiationInterface = "0.6, 0.7"
 Distributions = "0.25"
 JSON = "0.19 - 0.21, 1"
 LinearAlgebra = "<0.0.1, 1"
 MacroTools = "0.5"
 OrderedCollections = "1.8.1"
 Random = "1.6"
 StatsBase = "0.32, 0.33, 0.34"
+Test = "1"
 julia = "1.10.8"

ext/AbstractPPLDifferentiationInterfaceExt.jl

@@ -0,0 +1,100 @@
+module AbstractPPLDifferentiationInterfaceExt
+
+using AbstractPPL: AbstractPPL
+using AbstractPPL.Evaluators: Evaluators, Prepared, VectorEvaluator
+using ADTypes: AbstractADType, AutoReverseDiff
+using DifferentiationInterface: DifferentiationInterface as DI
+
+# Differentiate only `x`; the evaluator is passed as a `DI.Constant` context so
+# that in DynamicPPL the model and other evaluator state stay constant.
+@inline _call_evaluator(x, evaluator) = evaluator(x)
+
+struct DICache{F,GP,JP}
+    target::F
+    gradient_prep::GP
+    jacobian_prep::JP
+    use_context::Bool
+end
+
+# Compiled ReverseDiff only reuses a compiled tape on the one-argument path;
+# `DI.Constant` deactivates tape recording, so close the evaluator into the
+# target and call DI without contexts.
+function _prepare_di(prep::F, adtype::AutoReverseDiff{true}, x, evaluator) where {F}
+    target = Base.Fix2(_call_evaluator, evaluator)
+    return target, prep(target, adtype, x), false
+end
+
+function _prepare_di(prep::F, adtype::AbstractADType, x, evaluator) where {F}
+    return _call_evaluator, prep(_call_evaluator, adtype, x, DI.Constant(evaluator)), true
+end
+
+function AbstractPPL.prepare(
+    adtype::AbstractADType, problem, x::AbstractVector{<:Real}; check_dims::Bool=true
+)
+    evaluator = AbstractPPL.prepare(problem, x; check_dims)::VectorEvaluator
+    y = evaluator(x)
+    y isa Union{Number,AbstractVector} || throw(
+        ArgumentError(
+            "A prepared AD evaluator must return a scalar or AbstractVector; got $(typeof(y)).",
+        ),
+    )
+    if length(x) == 0
+        # DI prep crashes on length-0 input (e.g. ForwardDiff `BoundsError`); the
+        # `Val(0)` sentinel keeps the `gradient_prep === nothing` arity check meaningful.
+        gp, jp = y isa Number ? (Val(0), nothing) : (nothing, Val(0))
+        return Prepared(adtype, evaluator, DICache(_call_evaluator, gp, jp, true))
+    end
+    if y isa Number
+        target, gradient_prep, use_context = _prepare_di(
+            DI.prepare_gradient, adtype, x, evaluator
+        )
+        return Prepared(
+            adtype, evaluator, DICache(target, gradient_prep, nothing, use_context)
+        )
+    end
+    target, jacobian_prep, use_context = _prepare_di(
+        DI.prepare_jacobian, adtype, x, evaluator
+    )
+    return Prepared(adtype, evaluator, DICache(target, nothing, jacobian_prep, use_context))
+end
+
+@inline function AbstractPPL.value_and_gradient!!(
+    p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DICache}, x::AbstractVector{T}
+) where {T<:Real}
+    p.cache.gradient_prep === nothing &&
+        throw(ArgumentError("`value_and_gradient!!` requires a scalar-valued function."))
+    T <: Integer && Evaluators._reject_integer_input(x)
+    Evaluators._check_vector_length(p.evaluator.dim, x)
+    # Bypass DI on length-0 input — DI prep paths fail (e.g. ForwardDiff
+    # `BoundsError`); typed `T[]` matches the caller's element type.
+    length(x) == 0 && return (p.evaluator(x), T[])
+    return if p.cache.use_context
+        DI.value_and_gradient(
+            p.cache.target, p.cache.gradient_prep, p.adtype, x, DI.Constant(p.evaluator)
+        )
+    else
+        DI.value_and_gradient(p.cache.target, p.cache.gradient_prep, p.adtype, x)
+    end
+end
+
+@inline function AbstractPPL.value_and_jacobian!!(
+    p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DICache}, x::AbstractVector{T}
+) where {T<:Real}
+    p.cache.jacobian_prep === nothing &&
+        throw(ArgumentError("`value_and_jacobian!!` requires a vector-valued function."))
+    T <: Integer && Evaluators._reject_integer_input(x)
+    Evaluators._check_vector_length(p.evaluator.dim, x)
+    if length(x) == 0
+        val = p.evaluator(x)
+        return (val, similar(x, length(val), 0))
+    end
+    return if p.cache.use_context
+        DI.value_and_jacobian(
+            p.cache.target, p.cache.jacobian_prep, p.adtype, x, DI.Constant(p.evaluator)
+        )
+    else
+        DI.value_and_jacobian(p.cache.target, p.cache.jacobian_prep, p.adtype, x)
+    end
+end
+
+end # module