flax

c558a7c3 - Add test for multiple non final axes