code: def batch_quadratic_form(
      x: torch.Tensor, A: torch.Tensor
  ) -> torch.Tensor:
      """
      Compute the quadratic form x^T * A * x for a batched input x.
      The calcuation of pred_sigma (uncertainty) in LinUCB is done by quadratic form x^T * A^{-1} * x.
      Inspired by https://stackoverflow.com/questions/18541851/calculate-vt-a-v-for-a-matrix-of-vectors-v
      This is a vectorized implementation of out[i] = x[i].t() @ A @ x[i]
      x shape: (Batch, Feature_dim) or (Batch, Arm, Feature_dim)
      A shape: (Feature_dim, Feature_dim)
      output shape: (Batch) or (Batch, Arm)
      """
      return (torch.matmul(x, A) * x).sum(-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 2 for calling a global function with example usage: batch_quadratic_form(x, A) and returns: unknown