import jax.numpy as jnp
from flax import nnx
from jaxviz import trace_model



class MLP(nnx.Module):
    def __init__(self, rngs):
        self.dense0 = nnx.Linear(8, 16, rngs=rngs)
        self.dense1 = nnx.Linear(16, 4, rngs=rngs)

    def __call__(self, x):
        x = self.dense0(x)
        x = nnx.relu(x)
        x = self.dense1(x)
        return x


model = MLP(nnx.Rngs(0))
example_input = jnp.ones((1, 8))

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