class_docstr: Implements temperature scaling of logits. Based on results from "On Calibration of Modern Neural Networks":
  https://arxiv.org/abs/1706.04599. Temperature scaling scales all logits by the same constant factor. Though it
  may modify output probabilities it will never change argmax or categorical top-n predictions. In the case of
  binary classification with a threshold, however, calibration may change predictions.
  
  Implementation inspired by https://github.com/gpleiss/temperature_scaling
  
  Args:
      num_classes: The number of classes. Must be 2 if binary is True.
      binary: If binary is true, logits is expected to be a 1-dimensional array. If false, logits is a 2-dimensional
              array of shape (num_examples, num_classes).
class_docstr_tok: ["Implements","temperature","scaling","of","logits",".","Based","on","results","from","``","On","Calibration","of","Modern","Neural","Networks","''",":","https",":","\/\/arxiv.org\/abs\/1706.04599",".","Temperature","scaling","scales","all","logits","by","the","same","constant","factor",".","Though","it","may","modify","output","probabilities","it","will","never","change","argmax","or","categorical","top-n","predictions",".","In","the","case","of","binary","classification","with","a","threshold",",","however",",","calibration","may","change","predictions",".","Implementation","inspired","by","https",":","\/\/github.com\/gpleiss\/temperature_scaling","Args",":","num_classes",":","The","number","of","classes",".","Must","be","2","if","binary","is","True",".","binary",":","If","binary","is","true",",","logits","is","expected","to","be","a","1-dimensional","array",".","If","false",",","logits","is","a","2-dimensional","array","of","shape","(","num_examples",",","num_classes",")","."]
code: def train_calibration(
      self,
      logits: Union[torch.Tensor, np.ndarray],
      labels: Union[torch.Tensor, np.ndarray],
  ) -> CalibrationResult:
      logits = torch.as_tensor(
          logits, dtype=torch.float32, device=self.device
      )
      labels = torch.as_tensor(
          labels, dtype=torch.int64, device=self.device
      )
      one_hot_labels = nn.functional.one_hot(
          labels, self.num_classes
      ).float()
      if self.binary:
          # Treat binary classification as multi-class with 2 classes to re-use code.
          # The math works out the same: softmax([0, a])[1] == sigmoid(a)
          logits = torch.stack(
              [torch.zeros_like(logits), logits], axis=-1
          )
      nll_criterion = nn.CrossEntropyLoss().to(self.device)
      ece_criterion = ECELoss().to(self.device)
      # Saves the original temperature parameter, in case something goes wrong in optimization.
      original_temperature = self.temperature.clone().detach()
      self.temperature.requires_grad = True
      # Calculate NLL and ECE before temperature scaling
      before_calibration_nll = nll_criterion(
          logits, one_hot_labels
      ).item()
      before_calibration_ece = ece_criterion(
          logits, one_hot_labels
      ).item()
      logger.info(
          "Before temperature scaling:\n"
          "    Negative log-likelihood: %.3f\n"
          "    Expected Calibration Error: %.3f"
          % (before_calibration_nll, before_calibration_ece)
      )
  
      # Optimizes the temperature to minimize NLL
      optimizer = torch.optim.LBFGS(
          [self.temperature],
          lr=0.01,
          max_iter=50,
          line_search_fn="strong_wolfe",
      )
  
      def eval():
          optimizer.zero_grad()
          loss = nll_criterion(
              self.scale_logits(logits), one_hot_labels
          )
          loss.backward()
          return loss
  
      optimizer.step(eval)
  
      # Calculate NLL and ECE after temperature scaling
      after_calibration_nll = nll_criterion(
          self.scale_logits(logits), one_hot_labels
      ).item()
      after_calibration_ece = ece_criterion(
          self.scale_logits(logits), one_hot_labels
      ).item()
      logger.info("Optimal temperature: %.3f" % self.temperature.item())
      logger.info(
          "After temperature scaling:\n"
          "    Negative log-likelihood: %.3f\n"
          "    Expected Calibration Error: %.3f"
          % (after_calibration_nll, after_calibration_ece)
      )
      self.temperature.requires_grad = False
      # This should never happen, but if expected calibration error is higher after optimizing temperature, revert.
      if after_calibration_ece > before_calibration_ece:
          logger.warning(
              "Expected calibration error higher after scaling, "
              "reverting to temperature=%.3f."
              % original_temperature.item()
          )
          with torch.no_grad():
              self.temperature.data = original_temperature.data
      return CalibrationResult(
          before_calibration_nll,
          before_calibration_ece,
          after_calibration_nll,
          after_calibration_ece,
      )
  
imports: import logging
  from abc import ABC, abstractmethod
  from dataclasses import dataclass
  from typing import List, Type, Union
  import numpy
  import torch
  import torch.nn
  from ludwig.api_annotations import DeveloperAPI
  from ludwig.constants import BINARY, CATEGORY
  from ludwig.utils.registry import DEFAULT_KEYS, Registry
label_desc: Use image node_id 2 for calling the TemperatureScaling obj's underlying member method code with example usage: obj.train_calibration(logits, labels) and returns: CalibrationResult, loss