pmap changes the behaviour of a NumPyro sampler

[jeremiecoullon]

Hello!

I documented this behaviour in the NumPyro issue tracker. Essentially I found that in some cases, pmap(f)(x[None])[0] and f(x) do not produce the same result.

I was wondering if this behaviour is expected from pmap? And is it perhaps due to numerical precision errors?

[sharadmv]

It can vary from backend to backend (CPU vs GPU vs TPU) but in general, JIT compilation (including pmap) can change the numerics of a function due to compiler optimizations.

What backend are you using?

[jeremiecoullon]

Thanks for your answer, and sorry I never answered this!

I'm using CPU as backend. So in this context, changing the numerics of a function due to compiler optimizations can lead to numerical errors. Is that what your mean?