text: #!/usr/bin/env python3
  # Copyright (c) Meta Platforms, Inc. and affiliates.
  # All rights reserved.
  #
  # This source code is licensed under the BSD-style license found in the
  # LICENSE file in the root directory of this source tree.
  
  from typing import Any, cast, Dict, List, Optional, Type
  
  import torch
  from torchrec.metrics.metrics_namespace import MetricName, MetricNamespace, MetricPrefix
  from torchrec.metrics.rec_metric import (
      MetricComputationReport,
      RecMetric,
      RecMetricComputation,
      RecMetricException,
  )
  
  CTR_NUM = "ctr_num"
  CTR_DENOM = "ctr_denom"
  
  
  def compute_ctr(ctr_num: torch.Tensor, ctr_denom: torch.Tensor) -> torch.Tensor:
      return torch.where(ctr_denom == 0.0, 0.0, ctr_num / ctr_denom).double()
  
  
  def get_ctr_states(
      labels: torch.Tensor, predictions: torch.Tensor, weights: torch.Tensor
  ) -> Dict[str, torch.Tensor]:
      return {
          CTR_NUM: torch.sum(labels * weights, dim=-1),
          CTR_DENOM: torch.sum(weights, dim=-1),
      }
  
  
  class CTRMetricComputation(RecMetricComputation):
      r"""
      This class implementation the RecMetricComputation for CTR.
  
      The constructer arguments are defined in RecMetricComputation.
      See the docstring of RecMetricComputation for more detail.
      """
  
      def __init__(self, *args: Any, **kwargs: Any) -> None:
          super().__init__(*args, **kwargs)
          self._add_state(
              CTR_NUM,
              torch.zeros(self._n_tasks, dtype=torch.double),
              add_window_state=True,
              dist_reduce_fx="sum",
              persistent=True,
          )
          self._add_state(
              CTR_DENOM,
              torch.zeros(self._n_tasks, dtype=torch.double),
              add_window_state=True,
              dist_reduce_fx="sum",
              persistent=True,
          )
  
      def update(
          self,
          *,
          predictions: Optional[torch.Tensor],
          labels: torch.Tensor,
          weights: Optional[torch.Tensor]
      ) -> None:
          if predictions is None or weights is None:
              raise RecMetricException(
                  "Inputs 'predictions' and 'weights' should not be None for CTRMetricComputation update"
              )
          num_samples = predictions.shape[-1]
          for state_name, state_value in get_ctr_states(
              labels, predictions, weights
          ).items():
              state = getattr(self, state_name)
              state += state_value
              self._aggregate_window_state(state_name, state_value, num_samples)
  
      def _compute(self) -> List[MetricComputationReport]:
          return [
              MetricComputationReport(
                  name=MetricName.CTR,
                  metric_prefix=MetricPrefix.LIFETIME,
                  value=compute_ctr(
                      cast(torch.Tensor, self.ctr_num),
                      cast(torch.Tensor, self.ctr_denom),
                  ),
              ),
              MetricComputationReport(
                  name=MetricName.CTR,
                  metric_prefix=MetricPrefix.WINDOW,
                  value=compute_ctr(
                      self.get_window_state(CTR_NUM),
                      self.get_window_state(CTR_DENOM),
                  ),
              ),
          ]
  
  
  class CTRMetric(RecMetric):
      _namespace: MetricNamespace = MetricNamespace.CTR
      _computation_class: Type[RecMetricComputation] = CTRMetricComputation
  