class_docstr: A linear regression model for LinUCB.
  Note that instead of being trained by a PyTorch optimizer, we explicitly
      update attributes A and b (according to the LinUCB formulas implemented in
      reagent.training.cb.linucb_trainer.LinUCBTrainer).
  Since computing the regression coefficients inverse matrix inversion (expensive op), we
      save time by only computing the coefficients when necessary (when doing inference).
  
  Args:
      input_dim: Dimension of input data
      l2_reg_lambda: The weight on L2 regularization
      ucb_alpha: The coefficient on the standard deviation in UCB formula.
          Set it to 0 to predict the expected value instead of UCB.
      gamma: per-epoch discount factor (A and b get multiplied by gamma every epoch)
  
  Outputs:
  Dict {"pred_label": pred_label, "pred_sigma": pred_sigma, "ucb": ucb}
class_docstr_tok: ["A","linear","regression","model","for","LinUCB",".","Note","that","instead","of","being","trained","by","a","PyTorch","optimizer",",","we","explicitly","update","attributes","A","and","b","(","according","to","the","LinUCB","formulas","implemented","in","reagent.training.cb.linucb_trainer.LinUCBTrainer",")",".","Since","computing","the","regression","coefficients","inverse","matrix","inversion","(","expensive","op",")",",","we","save","time","by","only","computing","the","coefficients","when","necessary","(","when","doing","inference",")",".","Args",":","input_dim",":","Dimension","of","input","data","l2_reg_lambda",":","The","weight","on","L2","regularization","ucb_alpha",":","The","coefficient","on","the","standard","deviation","in","UCB","formula",".","Set","it","to","0","to","predict","the","expected","value","instead","of","UCB",".","gamma",":","per-epoch","discount","factor","(","A","and","b","get","multiplied","by","gamma","every","epoch",")","Outputs",":","Dict","{","``","pred_label","''",":","pred_label",",","``","pred_sigma","''",":","pred_sigma",",","``","ucb","''",":","ucb","}"]
code: def __init__(
      self,
      input_dim: int,
      *,
      l2_reg_lambda: float = 1.0,
      ucb_alpha: float = 1.0,
      gamma: float = 1.0,
  ) -> None:
      super().__init__(input_dim=input_dim)
  
      self.ucb_alpha = ucb_alpha
      self.l2_reg_lambda = l2_reg_lambda
      self.gamma = gamma
      assert self.gamma <= 1.0 and self.gamma > 0.0
  
      """
          the buffers below are split between "all data" and "current epoch" values. This is done
            to enable distributed training. "current epoch" values get reduced acorss all trainers at
            the end of an epoch (technically, whenever we estimate the coefficients, which could sometimes
            happen multiple times per epoch) and the sum gets redcuced with "all data" values and the new "all data"
            value is set to the reduction of "all data" and "current epoch" values.
          """
      # A is weighted average of X^T*X across all data
      self.register_buffer(
          "avg_A", torch.zeros(self.input_dim, self.input_dim)
      )
      # b is weighted average of label*X across all data
      self.register_buffer("avg_b", torch.zeros(self.input_dim))
      # A is weighted average of X^T*X across current epoch
      self.register_buffer(
          "cur_avg_A", torch.zeros(self.input_dim, self.input_dim)
      )
      # b is weighted average of label*X across current epoch
      self.register_buffer("cur_avg_b", torch.zeros(self.input_dim))
  
      self.register_buffer("_coefs", torch.zeros(self.input_dim))
      self.register_buffer(
          "inv_avg_A", torch.zeros(self.input_dim, self.input_dim)
      )
      self.register_buffer(
          "coefs_valid_for_avg_A",
          -torch.ones((self.input_dim, self.input_dim)),
      )  # value of avg_A matrix for which self.coefs were estimated
      self.register_buffer("num_obs", torch.zeros(1, dtype=torch.int64))
      self.register_buffer(
          "cur_num_obs", torch.zeros(1, dtype=torch.int64)
      )
      # initialize sum of weights below at small values to avoid dividing by 0
      self.register_buffer(
          "sum_weight", 1e-5 * torch.ones(1, dtype=torch.float)
      )
      self.register_buffer(
          "cur_sum_weight", 1e-5 * torch.ones(1, dtype=torch.float)
      )
  
      # add a dummy parameter so that DDP doesn't compain about lack of parameters with gradient
      self.dummy_param = torch.nn.parameter.Parameter(torch.zeros(1))
  
file_path: ReAgent/reagent/models/linear_regression.py
imports: import logging
  from typing import Dict, Optional
  import torch
  from pytorch_lightning.utilities.distributed import ReduceOp, sync_ddp_if_available
  from reagent.models.cb_base_model import UCBBaseModel
label_desc: Use image node_id 1 to create a new LinearRegressionUCB object from inherited base classes: UCBBaseModel with example: obj = LinearRegressionUCB(input_dim)