Massive PRNGKey allocations

[femtomc]

I'm running a large computation with scans, with vmap inside of scans, etc.

Before the main scan computation, I pre-allocate a set of sub_keys. However, I'm still running into OOM errors (massive allocations) for jax.random.split calls in the middle of my computation. The below is one of 8 (!) buffers. So massive OOM.

How do I avoid this? I'm doing lots of sampling, so it seems unavoidable to call jax.random.split -- on the other hand, I wish there was some way to evolve a set of keys in scan without newly allocating.

Also, what are best practices for keeping tracking of keys and reducing key allocations?

[froystig]

Can you make use of jax.random.fold_in within the scan to avoid upfront allocation via split? For example, if you increment an integer counter between each scan iteration, you can fold that in to a fixed "root" scan key to produce a local subkey for each scan iteration.

[femtomc]

Great tip! I'll try immediately.