Question to community: How important is being able to compile JAXNS into your JAX program?

[Joshuaalbert]

How important is being able to compile JAXNS into your JAX program?

Pros

1. You can embed it into a single compiled executable, which you can ship to CPU or GPU.
2. You get better performance (sometimes).

Cons

1. Worse memory management.
2. Distribution over compute resources is not great, as JAX relies on synchronous homogeneous SIMD (one single instruction multiple data kernel at a time over the same type of devices). And sampling best leverages concurrency, heterogeneous computing, and non-SIMD paradigm.
3. Performance is often worse than it should be because the algorithm doesn't nicely fit into pmap or shard_map. You always end up wasting wall-time.
4. Lose ability to make dynamic choices.
5. Much hard to distribute over massive compute clusters.

Alternative: What if the the likelihood is the only thing that needs to be JAX?

What you get:

1. you can still leverage auto-diff
2. you can still ship to CPU or GPU or TPU
3. you can distribute over massive compute clusters
4. you can easily stop and resume
5. almost no noticeable drop in performance
6. improved performance in many cases
7. but the algorithm itself is not pure JAX, so you cannot compile it.

[Joshuaalbert]

@tavin

[tavin]

@Joshuaalbert my only reason for compiling the jaxns call is performance. So it sounds like an interesting direction.

In any case in v2 I'm looking into compiling only the _main_ns_thread so that I can fall back into normal python and figure out how to refine my termination condition before I resume.

[Joshuaalbert]

@tavin Thanks for the input, I think performance will be better with this direction. It's still composed of compiled units of JAX code, but also makes use of it's own distributed fabric for executing over many devices/machines, as JAX's provided distribution functionality is very limited.