Skip to content

Let compile_train take extra updates - #76

Merged
jessegrabowski merged 2 commits into
pymc-devs:mainfrom
jessegrabowski:compile-extra-updates
Aug 10, 2026
Merged

Let compile_train take extra updates#76
jessegrabowski merged 2 commits into
pymc-devs:mainfrom
jessegrabowski:compile-extra-updates

Conversation

@jessegrabowski

Copy link
Copy Markdown
Member

compile_train couldn't write state that no gradient produces — a target-network sync, a Polyak average, replay priorities. The obvious workaround, compile_kwargs={"updates": ...}, hit a TypeError naming an internal function.

Now there's an extra_updates argument, and an updates key in compile_kwargs is taken as one. Either spelling raises if it writes a variable the step already writes — a target sync quietly overwriting an optimizer buffer would leave the rule configured but doing nothing.

Every state write no gradient produces was shut out of the ergonomic path -- target-network syncs, Polyak averages, replay priorities -- since the only opening collided with this function's own call to the compiler.
That is where a pytensor user looks first, and forwarding it raised a TypeError naming an internal function rather than saying anything useful.
@codecov-commenter

codecov-commenter commented Aug 10, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 95.42%. Comparing base (cf70a80) to head (7b3e308).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files
@@            Coverage Diff             @@
##             main      #76      +/-   ##
==========================================
+ Coverage   95.39%   95.42%   +0.02%     
==========================================
  Files          45       45              
  Lines        1782     1792      +10     
==========================================
+ Hits         1700     1710      +10     
  Misses         82       82              

☔ View full report in Codecov by Harness.
📢 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.

@jessegrabowski
jessegrabowski merged commit f60deb7 into pymc-devs:main Aug 10, 2026
7 checks passed
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.

2 participants