code: def test_mean_ensembler(
      self,
      multi_head=False,
      add_mean_last_layer_predictions=False,
      diff_last_layer_shapes=False,
  ):
      with context.graph_mode():
          ensembler = ensemble.MeanEnsembler(
              add_mean_last_layer_predictions=add_mean_last_layer_predictions
          )
          last_layer_dims = [3, 3]
          if diff_last_layer_shapes:
              last_layer_dims = [3, 5]
          if multi_head:
              subnetworks = [
                  self._build_subnetwork(
                      multi_head=multi_head,
                      last_layer_dim=last_layer_dim,
                  )
                  for last_layer_dim in last_layer_dims
              ]
          else:
              subnetworks = [
                  self._build_subnetwork(last_layer_dim=last_layer_dim)
                  for last_layer_dim in last_layer_dims
              ]
  
          if diff_last_layer_shapes:
              with self.assertRaisesRegexp(
                  ValueError,
                  r"Shape of \`last_layer\` tensors must be same",
              ):
                  built_ensemble = ensembler.build_ensemble(
                      subnetworks=subnetworks,
                      previous_ensemble_subnetworks=None,
                      features=None,
                      labels=None,
                      logits_dimension=None,
                      training=None,
                      iteration_step=None,
                      summary=None,
                      previous_ensemble=None,
                      previous_iteration_checkpoint=None,
                  )
              return
          built_ensemble = ensembler.build_ensemble(
              subnetworks=subnetworks,
              previous_ensemble_subnetworks=None,
              features=None,
              labels=None,
              logits_dimension=None,
              training=None,
              iteration_step=None,
              summary=None,
              previous_ensemble=None,
              previous_iteration_checkpoint=None,
          )
  
          with self.test_session() as sess:
              sess.run(tf_compat.v1.global_variables_initializer())
              got_logits = sess.run(built_ensemble.logits)
  
              if add_mean_last_layer_predictions:
                  got_predictions = sess.run(built_ensemble.predictions)
  
              logits = sess.run([s.logits for s in subnetworks])
              last_layer = sess.run([s.last_layer for s in subnetworks])
              if not multi_head:
                  expected_logits = np.mean(logits, axis=0)
                  expected_predictions = {
                      ensemble.MeanEnsemble.MEAN_LAST_LAYER: np.mean(
                          last_layer, axis=0
                      )
                  }
              else:
                  expected_logits = {
                      head_name: np.mean(
                          [s[head_name] for s in logits], axis=0
                      )
                      for head_name in multi_head
                  }
                  expected_predictions = {
                      "{}_{}".format(
                          ensemble.MeanEnsemble.MEAN_LAST_LAYER,
                          head_name,
                      ): np.mean(
                          [s[head_name] for s in last_layer], axis=0
                      )
                      for head_name in multi_head
                  }
  
              self.assertAllClose(expected_logits, got_logits)
              if add_mean_last_layer_predictions:
                  self.assertAllClose(
                      expected_predictions, got_predictions
                  )
  
file_path: adanet/adanet/ensemble/mean_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  from adanet import ensemble
  from adanet import subnetwork
  from adanet import tf_compat
  import numpy
  import tensorflow.compat.v2
  from tensorflow.python.eager import context
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 2 for calling the MeanTest obj's underlying member method code with example usage: obj.test_mean_ensembler(multi_head, add_mean_last_layer_predictions, diff_last_layer_shapes) without return types