myagues

You can use [`parameter_overview`](https://github.com/google/CommonLoopUtils/blob/fc344ef78f24758a96a1518500d08d0a5f01a6cd/clu/parameter_overview.py#L199-L231) in [`clu`](https://github.com/google/CommonLoopUtils) for params visualization: Commands ```python import jax import numpy as np from flax import linen as nn from clu import parameter_overview class CNN(nn.Module): @nn.compact...