Flax: Overview of Convolution and Sequence Models

Let’s look at the structures for building CNNs and RNNs.


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.

Get hands-on with 1200+ tech skills courses.