How to sumarize the model

[pisut-thengtubtim]

How can I sumarize the model like keras ?

Thank you.

[avital]

I think @rolandgvc maybe had some code that does this? Roland, has your summary printer code been ported to Linen?

[rolandgvc]

I think there was a clash of opinions after which the PR hasn't progressed much. @pisut-thengtubtim have a look at #288, it should be quite easy to make it work with the params dictionary.

[pisut-thengtubtim]

Thank you very much, it's working perfectly fine.

[andsteing]

Or you can use the clu library:

from clu import parameter_overview

parameter_overview.log_parameter_overview(params)
# alternatively:
# print(parameter_overview.get_parameter_overview(params))
# parameter_overview = parameter_overview.get_parameter_rows(params)

[pisut-thengtubtim]

I have tried this one also, it's working but I have to update the code at make_row(name, value) function
from
mean=float(value.mean()),
std=float(value.std()),
to
mean=value.mean().astype(float),
std=value.std().astype(float),

Thank you very much.

[andsteing]

Not sure why you have to do this. The following snippet works without modifications:

from clu import parameter_overview
from flax import linen as nn
import jax
import jax.numpy as jnp

rng = jax.random.PRNGKey(0)
x = jnp.zeros([5, 10])
dense = nn.Dense(7)
params = dense.init(rng, x)

print(parameter_overview.get_parameter_overview(params))