module (Any) – The linen module describing the architecture.
params (Any) – The variables produced by module.init(...) (e.g. {'params': ..., 'batch_stats':
...}).
sample_input (Any) – Required when the model contains convolution / spatial batch-norm layers.
