code: def __init__(self, *, rngs: nnx.Rngs):
      self.backbone = nnx.Linear(784, 128, rngs=nnx.Rngs(0))
      self.head = nnx.Linear(128, 10, rngs=rngs)
  
imports: import jax
  from flax.experimental import nnx
label_desc: Use image node_id 1 to create a new Classifier object from inherited base classes: nnx with example: obj = Classifier()