code: def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):
      self.count = Count(jnp.array(0))
      self.linear1 = Linear(din, dhidden, rngs=rngs)
      self.linear2 = Linear(dhidden, dout, rngs=rngs)
  
file_path: flax/flax/experimental/nnx/examples/toy_examples/02_lifted_transforms.py
imports: import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  from flax.experimental import nnx
label_desc: Use image node_id 1 to create a new MLP object from inherited base classes: nnx with example: obj = MLP(din, dhidden, dout)