Efficiency difference in using jax.lax.fori_loop vs looping over identical layers? about dm-haiku

[tomhennigan]

Hey @hrbigelow, both versions should work and in theory should be equally efficient, however we've seen a few cases (in particular with transformer models) where if you use structured control flow the XLA compiler does a better job at optimizing (in particular reducing peak memory usage) and (sometimes) overlapping communication with compute.

The recommended pattern in Haiku for repeated application of a block is to use hk.experimental.layer_stack.

The implementation of layer stack is kind of complex (it handles quite a few edge cases) but it basically boils down to using jax.lax.scan for the per-layer init and apply functions correctly.

[hrbigelow]

Thanks Tom. Actually I'm looking at the examples for hk.experimental.layer_stack. Just making sure I understand, it doesn't seem possible to somehow use layer_stack the same way you would use an ordinary hk.Module that has calls to hk.get_parameter, is that right?

Instead, if you wanted to build an hk.Module method that used layer_stack, you'd need to somehow obtain the pure function f to pass to stack(f)(...).