flax

7a429f33 - Remove the need for Conv and ConvTranspose to know the precise batch size.