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 the mean reward using the theta vectors.
  
  Args:
    theta: A `tf.Tensor` of theta vectors, shaped as [num_models,
      overall_context_dim].
    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 `tf.Tensor` of predicted mean reward shaped as
      [batch_size, self._num_actions].
docstr_tok: ["Predicts","the","mean","reward","using","the","theta","vectors",".","Args",":","theta",":","A","`","tf.Tensor","`","of","theta","vectors",",","shaped","as","[","num_models",",","overall_context_dim","]",".","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","`","tf.Tensor","`","of","predicted","mean","reward","shaped","as","[","batch_size",",","self._num_actions","]","."]
code: def _predict_mean_reward(
      self,
      theta: tf.Tensor,
      global_observation: tf.Tensor,
      arm_observations: Optional[tf.Tensor],
  ) -> tf.Tensor:
      """Predicts the mean reward using the theta vectors.
  
      Args:
        theta: A `tf.Tensor` of theta vectors, shaped as [num_models,
          overall_context_dim].
        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 `tf.Tensor` of predicted mean reward shaped as
          [batch_size, self._num_actions].
      """
      if self._accepts_per_arm_features:
          # The stacked observations are shaped as
          # [num_actions, batch_size, overall_context_dim].
          stacked_observations = tf.stack(
              [
                  self._get_current_observation(
                      global_observation, arm_observations, arm_index=k
                  )
                  for k in range(self._num_actions)
              ]
          )
          # theta is shaped as [1, overall_context_dim], because `num_models` is 1
          # for per-arm-features.
          # predicted_rewards is shaped [num_actions, batch_size, 1].
          predicted_rewards = tf.matmul(
              stacked_observations, theta, transpose_b=True
          )
          # Squeeze and transpose the predicted rewards to return a matrix of
          # shape [batch_size, num_actions].
          return tf.transpose(tf.squeeze(predicted_rewards, axis=-1))
      else:
          # global_observation: shaped [batch_size, global_context_dim].
          # self._theta: shaped [num_actions, global_context_dim].
          # predicted rewards shaped as [batch_size, num_actions].
          return tf.matmul(global_observation, theta, transpose_b=True)
  
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 3 for calling the LinearBanditPolicy obj's underlying member method code with example usage: obj._predict_mean_reward(theta, global_observation, arm_observations) and returns: tf, tf