//  [python]
def _distribution(self, time_step, policy_state):
    observation = time_step.observation
    if self.observation_and_action_constraint_splitter is not None:
      observation, _ = self.observation_and_action_constraint_splitter(
          observation
      )
    observation = tf.nest.map_structure(
        lambda o: t.cast(o, dtype=self._dtype), observation
    )
    global_observation, arm_observations = self._split_observation(observation)

    if self._add_bias:
      # The bias is added via a constant 1 feature.
      global_observation = tf.concat(
          [
              global_observation,
              tf.ones([tf.shape(global_observation)[0], 1], dtype=self._dtype),
          ],
          axis=1,
      )
    # Check the shape of the observation matrix. The observations can be
    # batched.
    if not global_observation.shape.is_compatible_with(
        [None, self._global_context_dim]
    ):
      raise ValueError(
          'Global observation shape is expected to be {}. Got {}.'.format(
              [None, self._global_context_dim],
              global_observation.shape.as_list(),
          )
      )
    global_observation = tf.reshape(
        global_observation, [-1, self._global_context_dim]
    )
    if (self._theta is not None) and (self._alpha < 1e-12):
      # When self._theta is provided and the exploration parameter is nearly
      # zero, inference only needs the predicted mean reward, which can be
      # computed efficiently by using the theta vectors.
      est_rewards = self._predict_mean_reward(
          self._theta, global_observation, arm_observations
      )
      rewards_for_argmax = est_rewards
    else:
      est_rewards, est_variances = self._predict_mean_reward_and_variance(
          global_observation, arm_observations
      )
      if self._exploration_strategy == ExplorationStrategy.optimistic:
        rewards_for_argmax = est_rewards + self._alpha * tf.sqrt(est_variances)
      elif self._exploration_strategy == ExplorationStrategy.sampling:
        mu_sampler = tfd.Normal(
            loc=est_rewards, scale=self._alpha * tf.sqrt(est_variances)
        )
        rewards_for_argmax = mu_sampler.sample()
      else:
        raise ValueError(
            'Exploraton strategy %s not implemented.'
            % self._exploration_strategy
        )

    mask = constraints.construct_mask_from_multiple_sources(
        time_step.observation,
        self._observation_and_action_constraint_splitter,
        (),
        self._num_actions,
    )
    if mask is not None:
      chosen_actions = policy_utilities.masked_argmax(
          rewards_for_argmax,
          mask,
          output_type=tf.nest.flatten(self._action_spec)[0].dtype,
      )
    else:
      chosen_actions = tf.argmax(
          rewards_for_argmax,
          axis=-1,
          output_type=tf.nest.flatten(self._action_spec)[0].dtype,
      )

    action_distributions = tfp.distributions.Deterministic(loc=chosen_actions)

    policy_info = policy_utilities.populate_policy_info(
        arm_observations,
        chosen_actions,
        rewards_for_argmax,
        est_rewards,
        self._emit_policy_info,
        self._accepts_per_arm_features,
    )

    return policy_step.PolicyStep(
        action_distributions, policy_state, policy_info
    )