code: def testNDCG(self, feature_key):
      # SubKeys will be a merger of top_k_list and sub_keys.
      metric = ndcg.NDCG(
          gain_key="gain", top_k_list=[1, 2]
      ).computations(
          sub_keys=[
              None,
              metric_types.SubKey(top_k=1),
              metric_types.SubKey(top_k=2),
          ],
          query_key="query",
          example_weighted=True,
      )[
          0
      ]
  
      query1_example1 = {
          "labels": np.array([1.0]),
          "predictions": np.array([0.2]),
          "example_weights": np.array([1.0]),
          feature_key: {
              "query": np.array(["query1"]),
              "gain": np.array([1.0]),
          },
      }
      query1_example2 = {
          "labels": np.array([0.0]),
          "predictions": np.array([0.8]),
          "example_weights": np.array([1.0]),
          feature_key: {
              "query": np.array(["query1"]),
              "gain": np.array([0.5]),
          },
      }
      query2_example1 = {
          "labels": np.array([0.0]),
          "predictions": np.array([0.5]),
          "example_weights": np.array([2.0]),
          feature_key: {
              "query": np.array(["query2"]),
              "gain": np.array([0.5]),
          },
      }
      query2_example2 = {
          "labels": np.array([1.0]),
          "predictions": np.array([0.9]),
          "example_weights": np.array([2.0]),
          feature_key: {
              "query": np.array(["query2"]),
              "gain": np.array([1.0]),
          },
      }
      query2_example3 = {
          "labels": np.array([0.0]),
          "predictions": np.array([0.1]),
          "example_weights": np.array([2.0]),
          feature_key: {
              "query": np.array(["query2"]),
              "gain": np.array([0.1]),
          },
      }
      query3_example1 = {
          "labels": np.array([1.0]),
          "predictions": np.array([0.9]),
          "example_weights": np.array([3.0]),
          feature_key: {
              "query": np.array(["query3"]),
              "gain": np.array([1.0]),
          },
      }
      query4_example1 = {
          "labels": np.array([1.0]),
          "predictions": np.array([0.9]),
          "example_weights": np.array([3.0]),
          feature_key: {
              "query": np.array(["query4"]),
              "gain": np.array([0.0]),  # 0 gain is ignored
          },
      }
      examples = [
          tfma_util.merge_extracts([query1_example1, query1_example2]),
          tfma_util.merge_extracts(
              [query2_example1, query2_example2, query2_example3]
          ),
          tfma_util.merge_extracts([query3_example1]),
          tfma_util.merge_extracts([query4_example1]),
      ]
  
      with beam.Pipeline() as pipeline:
          # pylint: disable=no-value-for-parameter
          result = (
              pipeline
              | "Create" >> beam.Create(examples)
              | "Process"
              >> beam.Map(metric_util.to_standard_metric_inputs, True)
              | "AddSlice" >> beam.Map(lambda x: ((), x))
              | "Combine" >> beam.CombinePerKey(metric.combiner)
          )
  
          # pylint: enable=no-value-for-parameter
  
          def check_result(got):
              try:
                  self.assertLen(got, 1)
                  got_slice_key, got_metrics = got[0]
                  self.assertEqual(got_slice_key, ())
                  ndcg1_key = metric_types.MetricKey(
                      name="ndcg",
                      sub_key=metric_types.SubKey(top_k=1),
                      example_weighted=True,
                  )
                  ndcg2_key = metric_types.MetricKey(
                      name="ndcg",
                      sub_key=metric_types.SubKey(top_k=2),
                      example_weighted=True,
                  )
                  # Query1 (weight=1): (p=0.8, g=0.5) (p=0.2, g=1.0)
                  # Query2 (weight=2): (p=0.9, g=1.0) (p=0.5, g=0.5) (p=0.1, g=0.1)
                  # Query3 (weight=3): (p=0.9, g=1.0)
                  #
                  # DCG@1:  0.5, 1.0, 1.0
                  # NDCG@1: 0.5, 1.0, 1.0
                  # Average NDCG@1: (1 * 0.5 + 2 * 1.0 + 3 * 1.0) / (1 + 2 + 3) ~ 0.92
                  #
                  # DCG@2: (0.5 + 1.0/log(3), (1.0 + 0.5/log(3), (1.0)
                  # NDCG@2: (0.5 + 1.0/log(3)) / (1.0 + 0.5/log(3)),
                  #         (1.0 + 0.5/log(3)) / (1.0 + 0.5/log(3)),
                  #         1.0
                  # Average NDCG@2: (1 * 0.860 + 2 * 1.0 + 3 * 1.0) / (1 + 2 + 3) ~ 0.97
                  self.assertDictElementsAlmostEqual(
                      got_metrics,
                      {ndcg1_key: 0.9166667, ndcg2_key: 0.9766198},
                      places=5,
                  )
  
              except AssertionError as err:
                  raise util.BeamAssertException(err)
  
          util.assert_that(result, check_result, label="result")
  
file_path: model-analysis/tensorflow_model_analysis/metrics/ndcg_test.py
imports: from absl.testing import parameterized
  import apache_beam
  from apache_beam.testing import util
  import numpy
  import tensorflow
  from tensorflow_model_analysis.eval_saved_model import testutil
  from tensorflow_model_analysis.metrics import metric_types
  from tensorflow_model_analysis.metrics import metric_util
  from tensorflow_model_analysis.metrics import ndcg
  from tensorflow_model_analysis.utils import util
label_desc: Use image node_id 1 for calling the NDCGMetricsTest obj's underlying member method code with example usage: obj.testNDCG(feature_key) without return types