code: def batch_quadratic_form_multi_arms(
      x: torch.Tensor, A: torch.Tensor
  ) -> torch.Tensor:
      """
      Compute the quadratic form x^T * A * x for a batched input x. (exploration term in UCB)
      x shape: (B, N)
      A shape: (num_arms, N, N)
      output shape: (B, num_arms)
      B is the batch size above
      N is the featur dimension here
      """
      # xta = x^T A
      # the input x is already transposed, so no transpose is applied below.
      # xta is dimension (num_arms, B, N)
      xta = torch.matmul(x, A)
      # einsum i: arm_idx; j: data index in batch; k: feature index
      return torch.einsum("ijk, jk -> ji", xta, x)
  
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 1 for calling a global function with example usage: batch_quadratic_form_multi_arms(x, A) and returns: torch