Error using batch with jit

[ravidziv]

I get the error Too many leaves for PyTreeDef; expected 6. when I'm trying to run the following code -

def get_network(W_std=1):
    init_fun, apply_fun, ker_fun = stax.serial(
        stax.Dense(1, W_std=W_std, b_std=0.1)
    )
    ker_fun =jit(batch(ker_fun, batch_size=25, device_count=0))
    kdd = ker_fun(train_xs, None)
    return 0
jit(get_network)(2.0)

[sschoenholz]

Thanks for bringing this up. I believe this issue should have been fixed by Roman's work on jit-compilation in batching. I've tried your repro in colab and it seems to work, but let me know if any problems persist.

[ravidziv]

Thanks, but it didn't solve it.
Colab code

[sschoenholz]

Ah I see, thanks! So this appears to work if you remove the jit around get_network, interesting! We'll look into it soon, but perhaps not until after ICLR.

One point that we should make more clear in the docs, especially when we have docs, is that we have found poor memory characteristics of applying jit to batch. We think this is an issue on the JAX / XLA end but haven't had time to pursue it to get them a simplified repro. If you run into OOM errors, it might be better to jit-before-batch, though we expect this situation to be temporary.

[ravidziv]

I understand, thanks!
Do Do you have an idea how can I vmap the batch for different networks?

[romanngg]

It looks like the latest JAX version has this fixed, and both jit and vmap work! Here's an adapted example from above: https://colab.research.google.com/gist/romanngg/ffdd9a41fdf5479eaeac95772c259d27/jit_or_vmap_of_network.ipynb

Please note that I am also unsure of how reasonable it is to apply jit to batch, especially while having a vmap on top. You may want to also try the options of
a) vmap / pmap on get_network, and have no batch or jit anywhere inside the function.
b) a simple python for loop over the W_stds of the jitted get_network, with no batch or jit inside, or with batch and no jit inside.
etc.
I suggest this because the purpose of batch is to sacrifice parallelization for the purpose of reducing memory cost. Therefore it seems redundant / sub-optimal to first call batch and then parallelize it with vmap or jit afterwards. But I'm no expert on how these JAX optimizations work so I very well might be wrong here...

[romanngg]

Closing, please let me know if there are still related issues!