code: def masked_argmax(
      input_tensor: types.Tensor,
      mask: types.Tensor,
      output_type: tf.DType = tf.int32,
  ) -> types.Tensor:
      """Computes the argmax where the allowed elements are given by a mask.
  
      If a row of `mask` contains all zeros, then this method will return -1 for the
      corresponding row of `input_tensor`.
  
      Args:
        input_tensor: Rank-2 Tensor of floats.
        mask: 0-1 valued Tensor of the same shape as input.
        output_type: Integer type of the output.
  
      Returns:
        A Tensor of rank 1 and type `output_type`, with the masked argmax of every
        row of `input_tensor`.
      """
      input_tensor.shape.assert_is_compatible_with(mask.shape)
      neg_inf = tf.constant(-float("Inf"), input_tensor.dtype)
      modified_input = tf.compat.v2.where(
          tf.cast(mask, tf.bool), input_tensor, neg_inf
      )
      argmax_tensor = tf.argmax(
          modified_input, axis=-1, output_type=output_type
      )
      # Replace results for invalid mask rows with -1.
      reduce_mask = tf.cast(tf.reduce_max(mask, axis=1), tf.bool)
      neg_one = tf.constant(-1, output_type)
      return tf.compat.v2.where(reduce_mask, argmax_tensor, neg_one)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  from typing import Optional, Sequence, Text, Union
  import tensorflow
  from tf_agents.specs import bandit_spec_utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 4 for calling a global function with example usage: masked_argmax(input_tensor, mask, output_type) and returns: tf