class_docstr: Implements matrix scaling of logits, as described in Beyond temperature scaling: Obtaining well-calibrated
  multiclass probabilities with Dirichlet calibration https://arxiv.org/abs/1910.12656.
  
  Unlike temperature scaling which has only one free parameter, matrix scaling has n_classes x (n_classes + 1)
  parameters. Use this only with a large validation set, as matrix scaling has a tendency to overfit small datasets.
  Also, unlike temperature scaling, matrix scaling can change the argmax or top-n predictions.
  
  NOTE: Matrix Scaling is not exposed in the UI or config yet, though it may be in a future release after testing.
  
  Args:
  num_classes: The number of classes.
  off_diagonal_l2: The regularization weight for off-diagonal matrix entries.
  mu: The regularization weight for bias vector. Defaults to off_diagonal_l2 if not specified.
class_docstr_tok: ["Implements","matrix","scaling","of","logits",",","as","described","in","Beyond","temperature","scaling",":","Obtaining","well-calibrated","multiclass","probabilities","with","Dirichlet","calibration","https",":","\/\/arxiv.org\/abs\/1910.12656",".","Unlike","temperature","scaling","which","has","only","one","free","parameter",",","matrix","scaling","has","n_classes","x","(","n_classes","+","1",")","parameters",".","Use","this","only","with","a","large","validation","set",",","as","matrix","scaling","has","a","tendency","to","overfit","small","datasets",".","Also",",","unlike","temperature","scaling",",","matrix","scaling","can","change","the","argmax","or","top-n","predictions",".","NOTE",":","Matrix","Scaling","is","not","exposed","in","the","UI","or","config","yet",",","though","it","may","be","in","a","future","release","after","testing",".","Args",":","num_classes",":","The","number","of","classes",".","off_diagonal_l2",":","The","regularization","weight","for","off-diagonal","matrix","entries",".","mu",":","The","regularization","weight","for","bias","vector",".","Defaults","to","off_diagonal_l2","if","not","specified","."]
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()
      nll_criterion = nn.CrossEntropyLoss().to(self.device)
      ece_criterion = ECELoss().to(self.device)
      self.w.requires_grad = True
      self.b.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 matrix scaling:\n"
          "    Negative log-likelihood: %.3f\n"
          "    Expected Calibration Error: %.3f"
          % (before_calibration_nll, before_calibration_ece)
      )
  
      # Optimizes the linear transform to minimize NLL
      optimizer = torch.optim.LBFGS(
          [self.w, self.b],
          lr=0.001,
          max_iter=200,
          line_search_fn="strong_wolfe",
      )
  
      def eval():
          optimizer.zero_grad()
          loss = (
              nll_criterion(self.scale_logits(logits), one_hot_labels)
              + self.regularization_terms()
          )
          loss.backward()
          return loss
  
      optimizer.step(eval)
  
      # Calculate NLL and ECE after matrix 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(
          "After matrix scaling:\n"
          "    Negative log-likelihood: %.3f\n"
          "    Expected Calibration Error: %.3f"
          % (after_calibration_nll, after_calibration_ece)
      )
      self.w.requires_grad = False
      self.b.requires_grad = False
      # This should never happen, but if expected calibration error is higher after optimizing matrix, revert.
      if after_calibration_ece > before_calibration_ece:
          logger.warning(
              "Expected calibration error higher after matrix scaling, reverting to identity."
          )
          with torch.no_grad():
              self.w.data = torch.eye(self.num_classes)
              self.b.data = torch.zeros(self.num_classes)
      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 MatrixScaling obj's underlying member method code with example usage: obj.train_calibration(logits, labels) and returns: CalibrationResult, loss