flax

7cb7c33e - Merge pull request #564 from freddyaboulton:linen-single-input-conv