parameters: (logits, legal_actions_mask)
return_statement: return exp_logit / tf.reduce_sum(exp_logit, axis=-1, keepdims=True)
docstring: Returns the softmax over the valid actions defined by `legal_actions_mask`.
  
    Args:
      logits: A tensor [..., num_actions] (e.g. [num_actions] or [B, num_actions])
        representing the logits to mask.
      legal_actions_mask: The legal action mask, same shape as logits. 1 means
        it's a legal action, 0 means it's illegal. If can be a tensorflow or numpy
        tensor.
docstring_summary: Returns the softmax over the valid actions defined by `legal_actions_mask`.
docstring_tokens: Returns the softmax over the valid actions defined by legal_actions_mask .
function: def tf_masked_softmax(logits, legal_actions_mask):
    """Returns the softmax over the valid actions defined by `legal_actions_mask`.
  
    Args:
      logits: A tensor [..., num_actions] (e.g. [num_actions] or [B, num_actions])
        representing the logits to mask.
      legal_actions_mask: The legal action mask, same shape as logits. 1 means
        it's a legal action, 0 means it's illegal. If can be a tensorflow or numpy
        tensor.
    """
    # This will raise a warning as we are taking the log of 0, which sets the 0
    # values to -inf. However, this is fine, as we then apply tf.exp, which sets
    # tf.exp(-inf) to 0. e.g. if we have logits [5, 3, 1], with legal_mask
    # [0, 1, 1], then masked_logits == [-inf, 3, 1], so we subtract the max to
    # get [-inf, 0, -2], and apply tf.exp to get [0, 1, e^-2].
    legal_actions_mask = tf.cast(legal_actions_mask, dtype=logits.dtype)
    masked_logits = logits + tf.log(legal_actions_mask)
    max_logit = tf.reduce_max(masked_logits, axis=-1, keepdims=True)
    exp_logit = tf.exp(masked_logits - max_logit)
    return exp_logit / tf.reduce_sum(exp_logit, axis=-1, keepdims=True)
function_tokens: def tf_masked_softmax ( logits , legal_actions_mask ) : # This will raise a warning as we are taking the log of 0, which sets the 0 # values to -inf. However, this is fine, as we then apply tf.exp, which sets # tf.exp(-inf) to 0. e.g. if we have logits [5, 3, 1], with legal_mask # [0, 1, 1], then masked_logits == [-inf, 3, 1], so we subtract the max to # get [-inf, 0, -2], and apply tf.exp to get [0, 1, e^-2]. legal_actions_mask = tf . cast ( legal_actions_mask , dtype = logits . dtype ) masked_logits = logits + tf . log ( legal_actions_mask ) max_logit = tf . reduce_max ( masked_logits , axis = - 1 , keepdims = True ) exp_logit = tf . exp ( masked_logits - max_logit ) return exp_logit / tf . reduce_sum ( exp_logit , axis = - 1 , keepdims = True )