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: Forward can return the mean or a UCB. If returning UCB, the CI width is stddev*ucb_alpha
  If ucb_alpha is not passed in, a fixed alpha from init is used
  
  inp: num_batch * dim
  self.coefs: num_arms * dim
  output: num_batch * num_arms
  output is in the format:
  [
      [score_1, score_2, ..., score_{num_arms}],
      ....,
  ]
code: def forward(
      self, inp: torch.Tensor, ucb_alpha: Optional[float] = None
  ) -> torch.Tensor:
      """
      Forward can return the mean or a UCB. If returning UCB, the CI width is stddev*ucb_alpha
      If ucb_alpha is not passed in, a fixed alpha from init is used
  
      inp: num_batch * dim
      self.coefs: num_arms * dim
      output: num_batch * num_arms
      output is in the format:
      [
          [score_1, score_2, ..., score_{num_arms}],
          ....,
      ]
      """
      if ucb_alpha is None:
          ucb_alpha = self.ucb_alpha
  
      results = torch.matmul(inp, self.coefs.t())
      if ucb_alpha == 0:
          return results
      results += ucb_alpha * torch.sqrt(
          batch_quadratic_form_multi_arms(inp, self.inv_A)
      )
      return results
  
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 4 for calling the DisjointLinearRegressionUCB obj's underlying member method code with example usage: obj.forward(inp, ucb_alpha) and returns: results, results