[nnx] add state summaries for print and display

[cgarciae]

What does this PR do?

• Adds parameter count and size per Variable type to both the penzai / nnx.display representation and to the string representation of NNX objects.
• Adds colors to the __str__ representation if the console supports it, the __repr__ version doesn't add any colors in case a simpler format is needed. This comes through improvements to nnx.reprlib.
• Slight simplification of the MNIST tutorial and Flax Basics.

from flax import nnx

class Block(nnx.Module):
  def __init__(self, din, dout, rngs: nnx.Rngs):
    self.linear = nnx.Linear(din, dout, rngs=rngs)
    self.bn = nnx.BatchNorm(dout, rngs=rngs)
    self.dropout = nnx.Dropout(0.2, rngs=rngs)

  def __call__(self, x):
    return nnx.relu(self.dropout(self.bn(self.linear(x))))

class Foo(nnx.Module):
  def __init__(self, rngs: nnx.Rngs):
    self.block1 = Block(32, 128, rngs=rngs)
    self.block2 = Block(128, 10, rngs=rngs)

  def __call__(self, x):
    return self.block2(self.block1(x))

foo = Foo(nnx.Rngs(0))

nnx.display(foo)
print(foo)

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB