Flax: Linear Modules
This lesson will focus on linear modules. These modules are helpful in designing feedforward neural networks. For simplicity, we’ll omit the optional arguments.
We can use
Dense() to apply an affine transformation to the previous layer, as:
Dense() applies a transformation only to the last dimension of the previous layer matrix, we can generalize it to any dimension by this function. For example, we can apply linear transformation along the required dimension through:
y = DenseGeneral(x,axis=<>)
Note: Most documentation incorrectly term the affine transformation as a linear transformation. This lesson will refer to them as affine thereon. (For further reading, check any linear algebra book such as Boyd’s Applied Linear Algebra (2018).)
If we print the output of a dense layer, some extra parameters will also be on display. Some of them are pretty helpful, like:
use_biascontrols whether or not to use the bias. It’s true by default.
kernel_initis the function to initialize the weight matrix,