import jax
import jax.numpy as jnp
import flax.linen as nn
from jaxviz import trace_model



class MLP(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(16)(x)
        x = nn.relu(x)
        x = nn.Dense(4)(x)
        return x


_mlp = MLP()
example_input = jnp.ones((1, 8))
_params = _mlp.init(jax.random.PRNGKey(0), example_input)


def model(x):
    return _mlp.apply(_params, x)

trace_model(model, example_input, view='per_device')