flax

15d556f9 - Merge pull request #3290 from chiamp:pjit_logical_partitioning