code: def test_head(
      self, head_fn, label_fn=None, loss_fn=None, metric_fn=None
  ):
      batch_size = 8
      input_shape = [batch_size, 32, 32, 3]
      hparams = hp.HParams(
          initial_architecture=[
              "FIXED_CHANNEL_CONVOLUTION_16",
              "FIXED_CHANNEL_CONVOLUTION_64",
              "PLATE_REDUCTION_FLATTEN",
          ],
          new_block_type="FIXED_CHANNEL_CONVOLUTION_32",
          learning_rate=1000.0,  # Approximating constant so we'll never diverge.
          optimizer="sgd",
      )
      instance = self._create_phoenix_instance(
          problem_type="cnn",
          input_shape=input_shape,
          head=head_fn(),
          loss_fn=loss_fn,
          metric_fn=metric_fn,
      )
      run_config = tf.estimator.RunConfig(
          model_dir=self.get_temp_dir() + "/1"
      )
  
      def input_fn():
          features = {"zeros": tf.zeros(input_shape)}
          if getattr(head_fn(), "_weight_column", None):
              features["weights"] = tf.ones(batch_size) * 0.5
          labels = (
              label_fn()
              if label_fn
              else tf.zeros(batch_size, dtype=tf.int32)
          )
          return features, labels
  
      estimator = instance.get_estimator(run_config, hparams, 10)
      estimator.train(input_fn=input_fn, max_steps=10)
      eval_result = estimator.evaluate(input_fn=input_fn, steps=10)
  
      self.assertAlmostEqual(0.0, eval_result["loss"], places=3)
  
imports: import os
  from absl import flags
  from absl.testing import parameterized
  from model_search import hparam
  from model_search import loss_fns
  from model_search import phoenix
  from model_search.architecture import architecture_utils
  from model_search.metadata import mock_metadata
  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 13 for calling the PhoenixTest obj's underlying member method code with example usage: obj.test_head(head_fn, label_fn, loss_fn, metric_fn) and returns: features,  labels