Part 1 of autograd backend: DifferentiationInterface extention (#158)

* Add DifferentiationInterface extension on top of evaluator interface

Layer the DifferentiationInterface AD extension onto the simplified evaluator
interface introduced on the `evaluators` branch. This is the catch-all path
for any `<:AbstractADType` backend without a native AbstractPPL extension.

- ext/AbstractPPLDifferentiationInterfaceExt.jl: DI-backed `prepare`,
  `value_and_gradient!!`, and `value_and_jacobian!!` for vector-input
  evaluators. The evaluator is passed as a `DI.Constant` so DynamicPPL-style
  problem state stays constant across calls. Compiled `AutoReverseDiff{true}`
  takes the one-argument tape path (the `DI.Constant` route would invalidate
  the compiled tape across calls). Empty inputs short-circuit before
  `DI.prepare_*` since many backends fail on length-zero arrays.
- Project.toml: DifferentiationInterface added as a weakdep with extension
  trigger and compat bound.
- test/autograd_tests.jl: shared problem definitions and `run_autograd_tests`
  entry point with separate gradient / jacobian / empty-input helpers.
- test/ext/differentiationinterface/: isolated test env using a local
  `DummyADType <: AbstractADType` to exercise the catch-all dispatch without
  pulling in a real AD package.
- .github/workflows/CI.yml: extend the ext matrix with
  `ext/differentiationinterface` alongside the existing
  `ext/logdensityproblems`.
- test/run_extras.jl: register the new label.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Inline AD output-shape assertions into the DI extension

The upstream `Simplify Evaluators module after review` commit dropped
`_assert_supported_output` / `_assert_jacobian_output` from
`Evaluators.jl`, expecting the AD-extension PRs that actually call them
to bring them back. Move both helpers into the DI extension, the sole
caller after rebase.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Fold AD output-shape checks into prepare's dispatch

Replace `_assert_supported_output` / `_assert_jacobian_output` with an
`if y isa Number / elseif y isa AbstractVector / else throw` cascade.
The jacobian assertion was unreachable after the scalar branch, and
the supported-output helper had a single call site; inlining the
remaining check as the `else` arm is clearer than two named helpers.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Move AD test cases into package test resources

Expose shared AD backend fixtures through TestResources so extension test environments can reuse the same cases without including files from test/.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* Use ForwardDiff in DI extension tests

Exercise the DifferentiationInterface catch-all with a real backend instead of a stub so the extension test covers the integration path used by downstream AD packages.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* Address PR review: simplify AD prepare and move test fixtures to Test extension

Replace `prepare(adtype, problem, x)`'s two-step `prepare` + `VectorEvaluator`
wrap with a single delegated call to `AbstractPPL.prepare(problem, x; check_dims)`,
eliminating a latent double-wrap when `check_dims=false`.

Move `TestResources` (callable problems, test cases, runner) out of `src/` and
into a new `AbstractPPLTestExt` package extension, so `Test` is loaded only when
downstream code opts in. Split the unified `TestCase` into `ValueCase` /
`ErrorCase`, drop the unused `:namedtuple` fixtures, and store the error
operation as a callable instead of a `:call`/`:gradient` symbol.

Flatten the DI extension test file to two `run_testcases(Val(...); ...)` calls.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Fix empty-input gradient arity, advertise LDP capabilities for DI

- value_and_gradient!! / value_and_jacobian!! now check arity before the
  length-zero short-circuit, so empty-input vector-valued / scalar-valued
  functions surface the correct ArgumentError instead of silently returning
  the wrong shape. Uses Val(0) as the cache sentinel for empty inputs.
- New AbstractPPLDifferentiationInterfaceLogDensityProblemsExt advertises
  LogDensityOrder{1} for any DI-prepared evaluator, registering the method
  in __init__ to side-step extension precompile-visibility constraints.
- Empty-input arity errors and the new capability advertisement are
  exercised in the DI test environment, which now also loads LDP.
- Project.toml [compat] alphabetised.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Unify LDP capabilities: Prepared = order 1, bare evaluator = order 0

Drop the triple extension and have AbstractPPLLogDensityProblemsExt
advertise LogDensityOrder{1} for any Prepared{<:Any,<:VectorEvaluator}.
The contract is now uniform: anything from `prepare(adtype, …)` claims
gradient capability; a bare VectorEvaluator from `prepare(problem, x)`
stays at order 0. Backends that don't implement value_and_gradient!! or
that wrap a vector-output function surface a runtime error at call time
rather than via a capability downgrade.

Documents the rule in docs/src/evaluators.md.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* format

* Skip Aqua persistent_tasks on Julia 1.10; drop dead --project from extras

Aqua's persistent_tasks check spawns a wrapper subprocess that runs
Pkg.precompile() on a package depending on AbstractPPL. On Julia 1.10
this hits "Declaring __precompile__(false) is not allowed in files
that are being precompiled" inside the wrapper's extension precompile
path. The dedicated Ext CI jobs load and exercise every extension on
min Julia and pass, so this is a Julia 1.10 / Aqua interaction, not a
defect in our extensions. Re-enable when min is bumped past 1.10.

The CI extras step also had a redundant --project=. that was overridden
by run_extras.jl's own Pkg.activate; drop it.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Polish evaluator API: public stubs, docstrings, AD-missing hints

- Mark generate_testcases / run_testcases as public (1.11+) so downstream
  AD-backend packages can reuse the conformance suite without reaching
  into private API.
- Note in the Prepared docstring that the two-arg constructor is for
  backends that allocate fresh storage per call.
- Expand the _ad_output_arity error message to call out matrix/tuple
  outputs explicitly and recommend flattening.
- Add a MethodError hint on value_and_gradient!! / value_and_jacobian!!
  that fires when no AD backend extension is loaded — also surfaces
  through LogDensityProblems.logdensity_and_gradient, which delegates
  to value_and_gradient!!. Suppressed once any backend registers a
  method, where the standard candidate list is more informative.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Update JULIA.md

* Address PR review: drop LDP ext, fix DI gradient/jacobian edges

- Drop AbstractPPLLogDensityProblemsExt and its test launcher; downstream
  packages can define LDP methods on `Prepared` directly when needed.
- Add `_check_vector_length` to enforce dim mismatch on
  `value_and_gradient!!`/`value_and_jacobian!!` before the `Val(0)`
  short-circuit and DI call (matches `VectorEvaluator{true}` behaviour).
- Unify gradient/jacobian DI prep into `_prepare_di(prep, ...)`; the
  `AutoReverseDiff{true}` carve-out now also fixes Jacobian tape recording.
- Tidy DI tests: remove LDP block, drop redundant constants and edge cases,
  tighten gradient-of-vector-valued case to `ArgumentError`.
- Document prep-time `evaluator(x)` call and reserved `Val{:vector}` /
  `Val{:edge}` testcase keys.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Reject integer inputs in DI value_and_{gradient,jacobian}!!

The bare `VectorEvaluator` rejected integer vectors with a clear
`ArgumentError`, but the DI derivative entry points went straight to
DI prep and surfaced an opaque `PreparationMismatchError`. Apply the
same compile-time `T <: Integer` check before dispatching to DI, and
add matching `:edge` testcases.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Add coverage upload to ext job, bump setup-julia and codecov-action versions

* Add regression test for AutoReverseDiff compiled tape no-context path

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Shravan Goswami <shravanngoswamii@gmail.com>

Author: 3 people

.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