pymc icon indicating copy to clipboard operation
pymc copied to clipboard

Write simple samplers in PyTensor (or Rust)

Open ricardoV94 opened this issue 1 month ago • 0 comments

Description

Samples with multiple steps like Slice could be written to loop in PyTensor. On Numba/JAX backends this should be much faster for models with cheap logp as we can skip Python altogether.

Simple-step samplers like Metropolis, when used alone (not blocked with other step samplers), could also benefit from this optimization.

Historically it didn't make much sense because Scan in the C-backend is pretty slow anyway. JAX and Numba fare much better, although there is room for improvement. Maybe only relevant after #7926

Could also write them in rust, which we know can interop well with numba as nutpie showed.

ricardoV94 avatar Oct 17 '25 11:10 ricardoV94