code: def __init__(self, hparams):
      self.hparams = hparams
  
      self._wpe = TrainVar(np.zeros([hparams.n_ctx, hparams.n_embd]))
      self._wte = TrainVar(np.zeros([hparams.n_vocab, hparams.n_embd]))
  
      self.blocks = objax.module.ModuleList(
          [
              Block(hparams.n_embd, hparams)
              for _ in range(hparams.n_layer)
          ]
      )
      self.norm = Norm(hparams.n_embd)
  
imports: import json
  import os
  import sys
  import jax.numpy
  import numpy
  import objax
  from objax.variable import TrainVar
label_desc: Use image node_id 1 to create a new Model object from inherited base classes: objax.module with example: obj = Model(hparams)