Seed the ADVI trainer's functions before compiling, as pymc does

[jessegrabowski]

Trainer now seeds its compiled step and sampling functions before compiling, following Model.compile_fn. It used to overwrite a compiled function's RNG storage afterward with jax_typify hardcoded, which worked on JAX and did nothing on MLX. A linker that copies RNGs at compile time now gets a fresh compile per seed. Every other linker gets one cached compile per model, reseeded in place with reseed_rngs.

The compiled functions are no longer attributes on the Trainer. It subclasses pymc's WithMemoization, and the unseeded compiles are locally_cachedmethod entries keyed on the model, so a fit against a new model compiles again and a fit against the same one does not.

The guide parameters and optimizer state now live in shared variables the Trainer creates once and hands to every compile. Each compiled step used to allocate its own Adam buffers, so a seeded recompile on JAX or MLX reset the moments mid-training. GradientTransformation gained a pytensor_init hook for those buffers.

The MLX case of test_reseeding_a_continued_fit is skipped. Compiling twice against one float64 shared variable stores a float32 mx.array into it, which is pymc-devs/pytensor#2378.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 93.24%. Comparing base (86fac3c) to head (1e6ae14).
⚠️ Report is 81 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #781       +/-   ##
===========================================
+ Coverage   51.60%   93.24%   +41.64%     
===========================================
  Files          73      108       +35     
  Lines        8003    10278     +2275     
===========================================
+ Hits         4130     9584     +5454     
+ Misses       3873      694     -3179     

Files with missing lines	Coverage Δ
pymc_extras/inference/advi/compile.py	97.82% <100.00%> (ø)
pymc_extras/inference/advi/optimizers.py	93.93% <100.00%> (ø)
pymc_extras/inference/advi/training.py	88.53% <100.00%> (+88.53%)	⬆️

... and 45 files with indirect coverage changes

🚀 New features to boost your workflow:

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