class_docstr: Class to build Bernoulli Thompson Sampling policies.
code: def _distribution(self, time_step, policy_state):
      if time_step.step_type.shape:
          if (
              tf.is_tensor(time_step.step_type)
              and time_step.step_type.shape.rank > 0
          ):
              batch_size = time_step.step_type.get_shape().as_list()[0]
          else:
              batch_size = 1
      else:
          batch_size = 1
      # Sample from the posterior distribution.
      posterior_dist = tfd.Beta(self._alpha, self._beta)
      predicted_reward_sampled = posterior_dist.sample([batch_size])
      predicted_reward_means_1d = tf.stack(
          [
              self._alpha[k] / (self._alpha[k] + self._beta[k])
              for k in range(self._expected_num_actions)
          ],
          axis=-1,
      )
      predicted_reward_means = tf.stack(
          [predicted_reward_means_1d for k in range(batch_size)], axis=0
      )
  
      mask = None
      if self._observation_and_action_constraint_splitter is not None:
          _, mask = self._observation_and_action_constraint_splitter(
              time_step.observation
          )
  
      # Argmax.
      if mask is not None:
          actions = policy_utilities.masked_argmax(
              predicted_reward_sampled,
              mask,
              output_type=self.action_spec.dtype,
          )
      else:
          actions = tf.argmax(
              predicted_reward_sampled,
              axis=-1,
              output_type=self.action_spec.dtype,
          )
  
      policy_info = policy_utilities.populate_policy_info(
          arm_observations=(),
          chosen_actions=actions,
          rewards_for_argmax=tf.cast(
              predicted_reward_sampled, tf.float32
          ),
          est_rewards=tf.cast(predicted_reward_means, tf.float32),
          emit_policy_info=self._emit_policy_info,
          accepts_per_arm_features=False,
      )
      if (
          policy_utilities.InfoFields.LOG_PROBABILITY
          in self._emit_policy_info
      ):
          policy_info._replace(
              log_probability=tf.zeros([batch_size], tf.float32)
          )
  
      return policy_step.PolicyStep(
          tfp.distributions.Deterministic(loc=actions),
          policy_state,
          policy_info,
      )
  
file_path: agents/tf_agents/bandits/policies/bernoulli_thompson_sampling_policy.py
imports: from typing import Optional, Sequence, Text
  import gin
  import tensorflow
  import tensorflow_probability
  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 BernoulliThompsonSamplingPolicy obj's underlying member method code with example usage: obj._distribution(time_step, policy_state) and returns: policy_step