flax

bf60167e - fix test_dense_general_batch_dim