class_docstr: A linear regression model for Disjoint LinUCB.
  Note that instead of being trained by a PyTorch optimizer, we explicitly
      update attributes A and b.
  
  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.
      gamma: The discount factor to avoid exploding numbers when doing incremental training
      using gamma as part of simplified D-LinUCB: https://arxiv.org/pdf/1909.09146.pdf
      In this simplified version of D-LinUCB, we calculate A = \sum \gamma^t xx^T
      to discount the old data. See N2441818 for why we do this.
      set gamma=1.0 to use traditional model
docstr: Compute current estimate of regression coefficients and A_inv=A**-1
  We save both coefficients and A_inv in case they are needed again before we add observations
  
  self.coefs: num_arms * dim
code: def _estimate_coefs(self):
      """
      Compute current estimate of regression coefficients and A_inv=A**-1
      We save both coefficients and A_inv in case they are needed again before we add observations
  
      self.coefs: num_arms * dim
      """
      self.A += sync_ddp_if_available(
          self.cur_A, reduce_op=ReduceOp.SUM
      )
      self.b += sync_ddp_if_available(
          self.cur_b, reduce_op=ReduceOp.SUM
      )
      # Set cur_A and cur_b back to 0 for next training epoch
      self.cur_A.zero_()
      self.cur_b.zero_()
      device = self.b.device
      # add regularization here so that it's not double-counted under distributed training
      # send them to the same device to avoid errors when doing dpp, example failed job without this: aienv-0d5c64a3b3
      m = self.A.to(device) + self.l2_reg_lambda * torch.eye(
          self.input_dim
      ).to(device)
      self.inv_A = torch.linalg.pinv(m).contiguous()
      assert self.inv_A.size()[0] == self.b.size()[0]
  
      # inv_A: (num_arms, d, d)
      # self.b: (num_arms, d)
      # einsum j: arm_idx, d: feature dimension index
      # output self.coefs: (num_arms, d)
      self.coefs = torch.einsum("jkl,jl->jk", self.inv_A, self.b)
      # copy A to make coefs_valid_for_A the same as A
      # need coefs_valid_for_A to check if we have done _estimate_coefs
      # this is needed to avoid redundant re-compute of coefs in forward function.
      self.coefs_valid_for_A = self.gamma * self.A.clone()
  
      cur_round_obs = sync_ddp_if_available(
          self.cur_num_obs.to(device), reduce_op=ReduceOp.SUM
      )
      logging.info(
          f"current round num of observations for {self.num_arms} arms are {cur_round_obs}"
      )
      logging.info(
          f"current round num of observations for all arms are {cur_round_obs.sum()}"
      )
      self.cur_num_obs.zero_()
  
file_path: ReAgent/reagent/models/disjoint_linucb_predictor.py
imports: import logging
  from typing import Optional
  import torch
  from pytorch_lightning.utilities.distributed import ReduceOp, sync_ddp_if_available
  from reagent.models.base import ModelBase
label_desc: Use image node_id 3 for calling the DisjointLinearRegressionUCB obj's underlying member method code with example usage: obj._estimate_coefs() without return types