class_docstr: Linear Bandit Policy to be used by LinUCB, LinTS and possibly others.
class_docstr_tok: ["Linear","Bandit","Policy","to","be","used","by","LinUCB",",","LinTS","and","possibly","others","."]
docstr: Predicts mean reward and variance.
  
  Args:
    global_observation: A `tf.Tensor` shaped as [batch_size,
      self._global_context_dim].
    arm_observations: An optional `tf.Tensor` shaped as [batch_size,
      self._num_actions, self._arm_context_dim]. Expected to be supplied only
      when self._accepts_per_arm_features is true.
  
  Returns:
    A size-2 tuple of `tf.Tensor`'s representing the predicted reward and
    variance, both shaped as [batch_size, self._num_actions].
docstr_tok: ["Predicts","mean","reward","and","variance",".","Args",":","global_observation",":","A","`","tf.Tensor","`","shaped","as","[","batch_size",",","self._global_context_dim","]",".","arm_observations",":","An","optional","`","tf.Tensor","`","shaped","as","[","batch_size",",","self._num_actions",",","self._arm_context_dim","]",".","Expected","to","be","supplied","only","when","self._accepts_per_arm_features","is","true",".","Returns",":","A","size-2","tuple","of","`","tf.Tensor","`","'s","representing","the","predicted","reward","and","variance",",","both","shaped","as","[","batch_size",",","self._num_actions","]","."]
code: def _predict_mean_reward_and_variance(
      self,
      global_observation: tf.Tensor,
      arm_observations: Optional[tf.Tensor],
  ) -> Tuple[tf.Tensor, tf.Tensor]:
      """Predicts mean reward and variance.
  
      Args:
        global_observation: A `tf.Tensor` shaped as [batch_size,
          self._global_context_dim].
        arm_observations: An optional `tf.Tensor` shaped as [batch_size,
          self._num_actions, self._arm_context_dim]. Expected to be supplied only
          when self._accepts_per_arm_features is true.
  
      Returns:
        A size-2 tuple of `tf.Tensor`'s representing the predicted reward and
        variance, both shaped as [batch_size, self._num_actions].
      """
      est_reward = []
      est_variance = []
      for k in range(self._num_actions):
          current_observation = tf.linalg.matrix_transpose(
              self._get_current_observation(
                  global_observation, arm_observations, k
              )
          )
          model_index = policy_utilities.get_model_index(
              k, self._accepts_per_arm_features
          )
          if self._use_eigendecomp:
              q_t_b = tf.matmul(
                  self._eig_matrix[model_index],
                  current_observation,
                  transpose_a=True,
              )
              lambda_inv = tf.divide(
                  tf.ones_like(self._eig_vals[model_index]),
                  self._eig_vals[model_index] + self._tikhonov_weight,
              )
              a_inv_x = tf.matmul(
                  self._eig_matrix[model_index],
                  tf.einsum("j,jk->jk", lambda_inv, q_t_b),
              )
          else:
              a_inv_x = linalg.conjugate_gradient(
                  self._cov_matrix[model_index]
                  + self._tikhonov_weight
                  * tf.eye(
                      self._overall_context_dim, dtype=self._dtype
                  ),
                  current_observation,
              )
          est_mean_reward = tf.einsum(
              "j,jk->k", self._data_vector[model_index], a_inv_x
          )
          est_reward.append(est_mean_reward)
          var = tf.einsum("ij,ij->j", current_observation, a_inv_x)
          est_variance.append(var)
  
      return (
          tf.stack(est_reward, axis=-1),
          tf.stack(est_variance, axis=-1),
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import enum
  from typing import Optional, Sequence, Text, Tuple
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.policies import constraints
  from tf_agents.bandits.policies import linalg
  from tf_agents.bandits.specs import utils
  from tf_agents.policies import tf_policy
  from tf_agents.policies import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.typing import types
label_desc: Use image node_id 4 for calling the LinearBanditPolicy obj's underlying member method code with example usage: obj._predict_mean_reward_and_variance(global_observation, arm_observations) and returns: tf,  tf