Skip to content

Attempt to make setindex!! on PA slices more understandable#1326

Draft
penelopeysm wants to merge 1 commit intomainfrom
py/refactor
Draft

Attempt to make setindex!! on PA slices more understandable#1326
penelopeysm wants to merge 1 commit intomainfrom
py/refactor

Conversation

@penelopeysm
Copy link
Member

This is Claude's partial attempt to refactor my ugly VNT code into smaller bits. Unfortunately, while the code here gives all the right results, it leads to type instabilities for less pathological cases, specifically setting x[1][1,2] when there's already an element inside. Consequently, this cannot be merged right now. I don't yet know if this can be worked around.

@github-actions
Copy link
Contributor

github-actions bot commented Mar 13, 2026

Benchmark Report

  • this PR's head: f883feb6714a6adee7c5075a8a7db1bcc9e62b5d
  • base branch: a31e3e396c1640b3344cf86041b2a7d9a52a74cf

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.89 │   346.26 │    0.81 │   8.56 │    6.65 │    1.29 │   2404.59 │   2302.22 │    1.04 │
│                   LDA │    12 │ reversediff │   true │  3150.54 │  2596.10 │    1.21 │   2.55 │    2.27 │    1.13 │   8046.76 │   5881.65 │    1.37 │
│   Loop univariate 10k │ 10000 │    mooncake │   true │ 31372.04 │ 32517.15 │    0.96 │   6.46 │    6.77 │    0.96 │ 202761.86 │ 219990.07 │    0.92 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│    Loop univariate 1k │  1000 │    mooncake │   true │  3179.00 │  3790.98 │    0.84 │   6.25 │    5.75 │    1.09 │  19859.67 │  21805.67 │    0.91 │
│      Multivariate 10k │ 10000 │    mooncake │   true │ 31893.11 │ 43586.48 │    0.73 │   9.96 │    7.50 │    1.33 │ 317563.59 │ 326907.91 │    0.97 │
│       Multivariate 1k │  1000 │    mooncake │   true │  3399.69 │  4767.86 │    0.71 │   9.38 │    6.76 │    1.39 │  31878.87 │  32239.56 │    0.99 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│ Simple assume observe │     1 │ forwarddiff │  false │     0.88 │     0.85 │    1.04 │  10.36 │   10.88 │    0.95 │      9.11 │      9.24 │    0.99 │
│           Smorgasbord │   201 │ forwarddiff │  false │   971.14 │  1122.81 │    0.86 │  72.79 │   79.69 │    0.91 │  70688.51 │  89475.88 │    0.79 │
│           Smorgasbord │   201 │      enzyme │   true │  1299.38 │  1452.11 │    0.89 │   6.47 │    4.50 │    1.44 │   8401.83 │   6530.06 │    1.29 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│           Smorgasbord │   201 │ forwarddiff │   true │  1306.87 │  1437.39 │    0.91 │  70.13 │   76.74 │    0.91 │  91646.20 │ 110305.82 │    0.83 │
│           Smorgasbord │   201 │    mooncake │   true │  1313.36 │  1452.46 │    0.90 │   4.60 │    3.91 │    1.18 │   6038.44 │   5673.50 │    1.06 │
│           Smorgasbord │   201 │ reversediff │   true │  1308.87 │  1432.49 │    0.91 │ 125.36 │  125.72 │    1.00 │ 164083.69 │ 180090.00 │    0.91 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│              Submodel │     1 │    mooncake │   true │     0.88 │     0.86 │    1.02 │  26.90 │   28.63 │    0.94 │     23.64 │     24.62 │    0.96 │
└───────────────────────┴───────┴─────────────┴────────┴──────────┴──────────┴─────────┴────────┴─────────┴─────────┴───────────┴───────────┴─────────┘

@codecov
Copy link

codecov bot commented Mar 13, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.34%. Comparing base (a31e3e3) to head (f883feb).

Additional details and impacted files
@@            Coverage Diff             @@
##             main    #1326      +/-   ##
==========================================
+ Coverage   78.26%   78.34%   +0.07%     
==========================================
  Files          50       50              
  Lines        3566     3583      +17     
==========================================
+ Hits         2791     2807      +16     
- Misses        775      776       +1     

☔ 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.

Comment on lines +672 to +686
# Step 2: Copy only the elements where value.mask is true.
# We mutate pa.data and pa.mask in place via views — this is safe because this
# function follows the `!!` convention (try to mutate, return the result).
# We do not do upfront type promotion here: the elements of `value` should have the
# same type as the elements of `pa`, since both originate from the same container.
data_view = view(pa.data, inds...; kw...)
mask_view = view(pa.mask, inds...; kw...)
for i in eachindex(value.mask)
if value.mask[i]
data_view[i] = value.data[i]
mask_view[i] = true
end
end

return pa
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude claimed there's no need to broaden the eltype of pa.data, but I think it's wrong; I think it does need to be broadened if necessary. The issue with that is that it breaks type stability for some tests.

@github-actions
Copy link
Contributor

DynamicPPL.jl documentation for PR #1326 is available at:
https://TuringLang.github.io/DynamicPPL.jl/previews/PR1326/

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant