flax

e1712906 - Allow linen's Conv layer to operate on arbitrary-rank inputs.

Allow linen's Conv layer to operate on arbitrary-rank inputs. This general solution handles both the case of multiple batch dimensions, and the (previously special-handled) case of an input without any batch dimension. PiperOrigin-RevId: 465512885