flax

c5dddd40 - Modifying linen Conv/Conv Transpose to accept int for kernel size and input without a batch dim.