make LogProbType a ref

[penelopeysm]

This is a less intrusive version of #1328, which allows the user to specify what base log-prob type they want. It is a bit more effort, but my hope is that it preserves type stability for 'normal' f64 usecases. I haven't yet tested this hypothesis.

[github-actions]

Benchmark Report

• this PR's head: 02d0d226f8d9de296b18042b11b4a11d5d81554e
• base branch: 34b8230ac30bb948798e58ec95adb65a9fbad4b4

Computer Information

Julia Version 1.11.9
Commit 53a02c0720c (2026-02-06 00:27 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 4 × AMD EPYC 7763 64-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, znver3)
Threads: 1 default, 0 interactive, 1 GC (on 4 virtual cores)

Benchmark Results

┌───────────────────────┬───────┬─────────────┬────────┬───────────────────────────────┬────────────────────────────┬─────────────────────────────────┐
│                       │       │             │        │       t(eval) / t(ref)        │     t(grad) / t(eval)      │        t(grad) / t(ref)         │
│                       │       │             │        │ ─────────┬──────────┬──────── │ ───────┬─────────┬──────── │ ──────────┬───────────┬──────── │
│                 Model │   Dim │  AD Backend │ Linked │     base │  this PR │ speedup │   base │ this PR │ speedup │      base │   this PR │ speedup │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│               Dynamic │    10 │    mooncake │   true │   280.45 │   272.43 │    1.03 │   8.14 │    7.60 │    1.07 │   2281.53 │   2070.95 │    1.10 │
│                   LDA │    12 │ reversediff │   true │  2601.34 │  3043.19 │    0.85 │   4.78 │    2.28 │    2.10 │  12422.54 │   6923.25 │    1.79 │
│   Loop univariate 10k │ 10000 │    mooncake │   true │ 31131.69 │ 29916.38 │    1.04 │   6.35 │    6.48 │    0.98 │ 197702.59 │ 193890.25 │    1.02 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│    Loop univariate 1k │  1000 │    mooncake │   true │  3895.06 │  3046.06 │    1.28 │   5.05 │    6.33 │    0.80 │  19685.68 │  19274.13 │    1.02 │
│      Multivariate 10k │ 10000 │    mooncake │   true │ 32069.13 │ 29908.48 │    1.07 │   9.84 │   10.12 │    0.97 │ 315454.40 │ 302711.55 │    1.04 │
│       Multivariate 1k │  1000 │    mooncake │   true │  3240.80 │  3228.65 │    1.00 │   9.79 │    9.44 │    1.04 │  31724.74 │  30478.01 │    1.04 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│ Simple assume observe │     1 │ forwarddiff │  false │     0.88 │   179.83 │    0.00 │  10.54 │    1.52 │    6.94 │      9.24 │    272.86 │    0.03 │
│           Smorgasbord │   201 │ forwarddiff │  false │   953.59 │ 20738.16 │    0.05 │  74.35 │   21.43 │    3.47 │  70897.81 │ 444464.78 │    0.16 │
│           Smorgasbord │   201 │      enzyme │   true │  1311.36 │  1245.41 │    1.05 │   4.75 │    4.82 │    0.98 │   6223.90 │   6008.57 │    1.04 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│           Smorgasbord │   201 │ forwarddiff │   true │  1283.95 │  1281.39 │    1.00 │  72.95 │   67.02 │    1.09 │  93660.91 │  85875.68 │    1.09 │
│           Smorgasbord │   201 │    mooncake │   true │  1321.87 │  1249.29 │    1.06 │   4.52 │    4.55 │    0.99 │   5975.32 │   5689.32 │    1.05 │
│           Smorgasbord │   201 │ reversediff │   true │  1315.89 │  1243.76 │    1.06 │ 121.90 │  125.29 │    0.97 │ 160407.39 │ 155828.11 │    1.03 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│              Submodel │     1 │    mooncake │   true │     0.88 │     0.84 │    1.04 │  26.96 │   26.78 │    1.01 │     23.69 │     22.55 │    1.05 │
└───────────────────────┴───────┴─────────────┴────────┴──────────┴──────────┴─────────┴────────┴─────────┴─────────┴───────────┴───────────┴─────────┘

[penelopeysm]

okay, logdensity and logdensity_and_gradient are fine, but this is even more type unstable 😳

DynamicPPL.LogProbType[] = Float64
@model f64() = x ~ Normal(0.0, 1.0)
benchmark(f64())
@inferred logjoint(f64(), (; x = randn()))
# ERROR: return type Float64 does not match inferred return type Any

[codecov]

Codecov Report

❌ Patch coverage is 75.00000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.17%. Comparing base (34b8230) to head (02d0d22).

Files with missing lines	Patch %	Lines
src/distribution_wrappers.jl	33.33%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1336      +/-   ##
==========================================
- Coverage   78.26%   74.17%   -4.10%     
==========================================
  Files          50       50              
  Lines        3566     3551      -15     
==========================================
- Hits         2791     2634     -157     
- Misses        775      917     +142     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.