Flax: Overview of Convolution and Sequence Models

Convolutions

The importance of convolutional neural networks (CNNs) in different applications of computer vision, and even NLP, is well-established.

Since we have already covered the mechanics of convolution in the earlier chapters, it should be sufficient just to show the respective Flax syntax. Remember that XLA only provides the main convolution functions while Flax provides the high-level wrappers. Wrappers for Flax are:

• Conv, used for traditional convolution.
• ConvTranspose, used for transposed convolution.