code: def _define_model(
      self,
      model_type,
      data_dim,
      num_hidden,
      kernel_initializer="ones",
      use_weight_norm=True,
      data_init=True,
  ):
      if use_weight_norm:
          base_layer = tfkl.Dense(
              num_hidden, kernel_initializer=kernel_initializer
          )
          layer = weight_norm.WeightNorm(
              base_layer, data_init=data_init, name="maybe_norm_layer"
          )
      else:
          layer = tfkl.Dense(
              num_hidden,
              kernel_initializer=kernel_initializer,
              name="maybe_norm_layer",
          )
  
      if model_type == "layer":
          return layer
      elif model_type == "sequential":
          return tfk.Sequential(
              [
                  tfkl.InputLayer((data_dim,)),
                  layer,
                  tfkl.Dense(1, kernel_initializer=kernel_initializer),
              ]
          )
      elif model_type == "sequential_no_input":
          return tfk.Sequential(
              [
                  layer,
                  tfkl.Dense(1, kernel_initializer=kernel_initializer),
              ]
          )
      elif model_type == "functional":
          inputs = tfkl.Input(shape=(data_dim,))
          net = layer(inputs)
          outputs = tfkl.Dense(
              1, kernel_initializer=kernel_initializer
          )(net)
          return tfk.Model(inputs=inputs, outputs=outputs)
      else:
          raise ValueError(
              "{} is not a valid model type".format(model_type)
          )
  
imports: import tempfile
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.layers import weight_norm
label_desc: Use image node_id 2 for calling the WeightNormTest obj's underlying member method code with example usage: obj._define_model(model_type, data_dim, num_hidden, kernel_initializer, use_weight_norm, data_init) and returns: layer, tfk, tfk, tfk