code: def _find_action_probabilities(
      greedy_action_prob: types.Tensor,
      other_actions_probs: types.Tensor,
      max_exploration_prob: float,
  ):
      """Finds action probabilities satisfying `max_exploration_prob`.
  
      Given action probabilities calculated by different values of the gamma
      parameter, this function attempts to find action probabilities at a specific
      gamma value such that non-greedy actions are chosen with at most
      `max_exploration_prob` probability. If such an upper bound can be achieved,
      the return maximizes the exploration probability subject to the upper bound.
      Otherwise, it minimizes the exploration probability on a best-effort basis.
  
      Args:
        greedy_action_prob: A tensor shaped as [batch_size, d], the probabilities of
          choosing the greedy action under `d` different values of gamma.
        other_actions_probs: A tensor shaped as [batch_size, num_actions, d], all
          non-greedy action probabilities under `d` different values of gamma. The
          last dimension is assumed to be aligned with that of `greedy_action_prob`.
        max_exploration_prob: A float, the maximum probability of choosing
          non-greedy actions.
  
      Returns:
        A tuple of two tensors for the greedy action probability and non-greedy
        actions probabilities, shaped as [batch_size, 1] and
        [batch_size, num_actions], respectively.
      """
      if greedy_action_prob.shape.rank != 2:
          raise ValueError(
              "`greedy_action_prob` is expected to be rank-2, but found otherwise:"
              f" {greedy_action_prob}"
          )
      if other_actions_probs.shape.rank != 3:
          raise ValueError(
              "`other_actions_probs` is expected to be rank-3, but found otherwise:"
              f" {other_actions_probs}"
          )
      if greedy_action_prob.shape[-1] != other_actions_probs.shape[-1]:
          raise ValueError(
              "`greedy_action_prob` and `other_actions_probs` are "
              "expected to have the same last dimension, but found "
              f"otherwise. `greedy_action_prob`: {greedy_action_prob}"
              f", `other_actions_probs`: {other_actions_probs}"
          )
  
      # A [batch_size, d] bool tensor indicating which elements of
      # `greedy_action_prob` satisfy the `max_exploration_prob` constraint.
      valid_gamma_mask = tf.greater_equal(
          greedy_action_prob, 1.0 - max_exploration_prob
      )
      # A [batch_size] bool tensor indicating the batch members that have at least
      # one valid entry in `greedy_action_prob`.
      feasible = tf.greater(
          tf.reduce_sum(tf.cast(valid_gamma_mask, tf.float32), axis=1),
          0.0,
      )
      # We mask the probability entries corresponding to invalid gamma values by 2.0
      # so that they will not be selected as minimizers. See further details in the
      # comment below.
      greedy_action_prob_masked = tf.where(
          valid_gamma_mask,
          greedy_action_prob,
          2.0 * tf.ones_like(greedy_action_prob),
      )
      # For batch members where the `max_exploration_prob` constraint is feasible,
      # we maximize the exploration probability (or equivalently, minimize the
      # greedy action probability) subject to the constraint via masking.
      # For batch members where the `max_exploration_prob` constraint is infeasible,
      # we simply minimize the exploration probability (or equivalently, maximize
      # the greedy action probability).
      gamma_indices = tf.where(
          feasible,
          tf.argmin(greedy_action_prob_masked, axis=1),
          tf.argmax(greedy_action_prob, axis=1),
      )
      gamma_indices = tf.expand_dims(gamma_indices, axis=-1)
      greedy_action_prob = tf.gather(
          greedy_action_prob, gamma_indices, axis=1, batch_dims=1
      )
      num_actions = tf.shape(other_actions_probs)[1]
      other_actions_probs = tf.gather(
          other_actions_probs,
          tf.tile(gamma_indices, [1, num_actions]),
          axis=2,
          batch_dims=2,
      )
      return greedy_action_prob, other_actions_probs
  
file_path: agents/tf_agents/bandits/policies/falcon_reward_prediction_policy.py
imports: from typing import Iterable, Optional, Sequence, Text, Tuple
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.policies import constraints
  from tf_agents.bandits.policies import reward_prediction_base_policy
  from tf_agents.distributions import shifted_categorical
  from tf_agents.policies import utils
  from tf_agents.typing import types
label_desc: Use image node_id 2 for calling a global function with example usage: _find_action_probabilities(greedy_action_prob, other_actions_probs, max_exploration_prob) and returns: greedy_action_prob,  other_actions_probs