Pjit and partitioning utilities

[joaogui1]

Will Haiku add partitioning utilities like Flax?

[tomhennigan]

Thanks for sharing, that's the first time I've seen those. We don't have any plans to do this at the moment, we have quite a few people using Haiku + pjit internally without needing additional utilities from Haiku.

Are there particular utilities that you think would be beneficial for working with pjit and Haiku?

[rosikand]

we have quite a few people using Haiku + pjit internally without needing additional utilities from Haiku.

Do you know of any public examples that use pjit for parallel NN training using Haiku?

[nlsfnr]

@rosikand It might be worth having a look at the Mesh Transformer JAX library, more concretely here