code: def test_weight_feature(self, is_multitask, weight_is_a_feature):
      # Force graph mode
      with tf.compat.v1.Graph().as_default():
          learning_rate_spec = {
              "learning_rate": 0.001,
              "gradient_max_norm": 3,
          }
          spec = phoenix_spec_pb2.PhoenixSpec(
              problem_type=phoenix_spec_pb2.PhoenixSpec.DNN
          )
          labels = tf.ones([20], dtype=tf.int32)
          if is_multitask:
              text_format.Merge(
                  """
              multi_task_spec {
                label_name: "label1"
                number_of_classes: 10
                weight_feature_name: "weight1"
                weight_is_a_feature: %s
              }
              multi_task_spec {
                label_name: "label2"
                number_of_classes: 10
                weight_feature_name: "weight2"
                weight_is_a_feature: %s
              }
          """
                  % (
                      str(weight_is_a_feature),
                      str(weight_is_a_feature),
                  ),
                  spec,
              )
              labels = {
                  "label1": tf.ones([20], dtype=tf.int32),
                  "label2": tf.ones([20], dtype=tf.int32),
              }
  
          weights = {
              "weight1": tf.constant([2] * 20),
              "weight2": tf.constant([3] * 20),
          }
          features = {"x": tf.zeros([10, 10])}
          if weight_is_a_feature:
              features.update(weights)
          elif isinstance(labels, dict):
              labels.update(weights)
          task_manager_instance = task_manager.TaskManager(
              spec,
              logits_dimension=None,
              loss_fn=loss_fns.make_multi_class_loss_fn(),
              head=None,
          )
          logits = tf.keras.layers.Dense(10)(tf.zeros([20, 10]))
          logits_spec = architecture_utils.LogitsSpec(logits=logits)
          fake_tower = collections.namedtuple(
              "fake_tower", ["logits_spec", "previous_model_dir"]
          )
          towers = {"search_generator": [fake_tower(logits_spec, None)]}
  
          _ = task_manager_instance.create_model_spec(
              features=features,
              params=hp.HParams(optimizer="sgd"),
              learning_rate_spec=learning_rate_spec,
              towers=towers,
              labels=labels,
              model_directory=self.get_temp_dir(),
              mode=tf.estimator.ModeKeys.TRAIN,
              my_id=1,
              trial_mode=trial_utils.TrialMode.NO_PRIOR,
              use_tpu=False,
              predictions_fn=_default_predictions_fn,
          )
  
imports: import collections
  from absl.testing import parameterized
  from model_search import hparam
  from model_search import loss_fns
  from model_search import task_manager
  from model_search.architecture import architecture_utils
  from model_search.generators import trial_utils
  from model_search.proto import phoenix_spec_pb2
  import numpy
  import tensorflow.compat.v2
  from google.protobuf import text_format
label_desc: Use image node_id 6 for calling the TaskManagerTest obj's underlying member method code with example usage: obj.test_weight_feature(is_multitask, weight_is_a_feature) without return types