code: def __init__(self, nin, nout):
      self.v1 = objax.TrainVar(jn.ones((nin, nout)))
      self.var_list = [
          objax.TrainVar(jn.zeros(nin)),
          objax.TrainVar(jn.zeros(nout)),
      ]
  
imports: import functools
  import inspect
  import unittest
  from typing import Tuple
  import numpy
  from jax import numpy
  import objax
  from objax.typing import JaxArray
label_desc: Use image node_id 1 to create a new SimpleModule object from inherited base classes: objax with example: obj = SimpleModule(nin, nout)