parameters: (
      inputs, init_h, kernel, recurrent_kernel, bias, mask, time_major,
      go_backwards, activation, recurrent_activation)
return_statement: return last_output, outputs, new_h, runtime
docstring: Call the GRU with optimized backend kernel selection.
  
    Under the hood, this function will create two TF function, one with the most
    generic kernel and can run on all device condition, and the second one with
    CuDNN specific kernel, which can only run on GPU.
  
    The first function will be called with normal_lstm_params, while the second
    function is not called, but only registered in the graph. The Grappler will
    do the proper graph rewrite and swap the optimized TF function based on the
    device placement.
  
    Args:
      inputs: Input tensor of GRU layer.
      init_h: Initial state tensor for the cell output.
      kernel: Weights for cell kernel.
      recurrent_kernel: Weights for cell recurrent kernel.
      bias: Weights for cell kernel bias and recurrent bias. Only recurrent bias
        is used in this case.
      mask: Boolean tensor for mask out the steps within sequence.
      time_major: Boolean, whether the inputs are in the format of
        [time, batch, feature] or [batch, time, feature].
      go_backwards: Boolean (default False). If True, process the input sequence
        backwards and return the reversed sequence.
      activation: Activation function to use for output.
      recurrent_activation: Activation function to use for hidden recurrent state.
  
    Returns:
      List of output tensors, same as standard_gru.
docstring_summary: Call the GRU with optimized backend kernel selection.
docstring_tokens: Call the GRU with optimized backend kernel selection .
function: def gru_with_backend_selection(
      inputs, init_h, kernel, recurrent_kernel, bias, mask, time_major,
      go_backwards, activation, recurrent_activation):
    """Call the GRU with optimized backend kernel selection.
  
    Under the hood, this function will create two TF function, one with the most
    generic kernel and can run on all device condition, and the second one with
    CuDNN specific kernel, which can only run on GPU.
  
    The first function will be called with normal_lstm_params, while the second
    function is not called, but only registered in the graph. The Grappler will
    do the proper graph rewrite and swap the optimized TF function based on the
    device placement.
  
    Args:
      inputs: Input tensor of GRU layer.
      init_h: Initial state tensor for the cell output.
      kernel: Weights for cell kernel.
      recurrent_kernel: Weights for cell recurrent kernel.
      bias: Weights for cell kernel bias and recurrent bias. Only recurrent bias
        is used in this case.
      mask: Boolean tensor for mask out the steps within sequence.
      time_major: Boolean, whether the inputs are in the format of
        [time, batch, feature] or [batch, time, feature].
      go_backwards: Boolean (default False). If True, process the input sequence
        backwards and return the reversed sequence.
      activation: Activation function to use for output.
      recurrent_activation: Activation function to use for hidden recurrent state.
  
    Returns:
      List of output tensors, same as standard_gru.
    """
    params = {
        'inputs': inputs,
        'init_h': init_h,
        'kernel': kernel,
        'recurrent_kernel': recurrent_kernel,
        'bias': bias,
        'mask': mask,
        'time_major': time_major,
        'go_backwards': go_backwards,
        'activation': activation,
        'recurrent_activation': recurrent_activation
    }
  
    def cudnn_gru_with_fallback(inputs, init_h, kernel, recurrent_kernel,
                                bias, mask, time_major, go_backwards, activation,
                                recurrent_activation):
      """Use CuDNN kernel when mask is none or strictly right padded."""
      if mask is None:
        return cudnn_gru(inputs=inputs, init_h=init_h, kernel=kernel,
                         recurrent_kernel=recurrent_kernel, bias=bias, mask=mask,
                         time_major=time_major, go_backwards=go_backwards)
      # Note that mask is a boolean tensor, which doesn't need to do gradient
      # calculation, when using tf.cond, a default gradient is added for it,
      # which then cause the backward function to have a signature mismatch.
      # Force the mask to not generate gradient to allow implementation_selector
      # to work properly.
      # TODO(b/80444525): Remove the stop_gradient().
      mask = array_ops.stop_gradient(mask)
  
      def input_right_padded():
        return cudnn_gru(inputs=inputs, init_h=init_h, kernel=kernel,
                         recurrent_kernel=recurrent_kernel, bias=bias, mask=mask,
                         time_major=time_major, go_backwards=go_backwards)
  
      def input_not_right_padded():
        return standard_gru(inputs=inputs, init_h=init_h, kernel=kernel,
                            recurrent_kernel=recurrent_kernel, bias=bias,
                            mask=mask, time_major=time_major,
                            go_backwards=go_backwards, activation=activation,
                            recurrent_activation=recurrent_activation)
  
      return control_flow_ops.cond(
          is_sequence_right_padded(mask, time_major),
          true_fn=input_right_padded,
          false_fn=input_not_right_padded)
  
    # Each time a `tf.function` is called, we will give it a unique
    # identifiable API name, so that Grappler won't get confused when it
    # sees multiple GRU layers added into same graph, and it will be able
    # to pair up the different implementations across them.
    api_name = 'gru_' + str(uuid.uuid4())
    defun_standard_gru = _generate_defun_backend(
        api_name, _CPU_DEVICE_NAME, standard_gru)
    defun_cudnn_gru = _generate_defun_backend(
        api_name, _GPU_DEVICE_NAME, cudnn_gru_with_fallback)
  
    # Call the normal GRU impl and register the CuDNN impl function. The
    # grappler will kick in during session execution to optimize the graph.
    last_output, outputs, new_h, runtime = defun_standard_gru(**params)
    function.register(defun_cudnn_gru, **params)
    return last_output, outputs, new_h, runtime
function_tokens: def gru_with_backend_selection ( inputs , init_h , kernel , recurrent_kernel , bias , mask , time_major , go_backwards , activation , recurrent_activation ) : params = { 'inputs' : inputs , 'init_h' : init_h , 'kernel' : kernel , 'recurrent_kernel' : recurrent_kernel , 'bias' : bias , 'mask' : mask , 'time_major' : time_major , 'go_backwards' : go_backwards , 'activation' : activation , 'recurrent_activation' : recurrent_activation } def cudnn_gru_with_fallback ( inputs , init_h , kernel , recurrent_kernel , bias , mask , time_major , go_backwards , activation , recurrent_activation ) : """Use CuDNN kernel when mask is none or strictly right padded.""" if mask is None : return cudnn_gru ( inputs = inputs , init_h = init_h , kernel = kernel , recurrent_kernel = recurrent_kernel , bias = bias , mask = mask , time_major = time_major , go_backwards = go_backwards ) # Note that mask is a boolean tensor, which doesn't need to do gradient # calculation, when using tf.cond, a default gradient is added for it, # which then cause the backward function to have a signature mismatch. # Force the mask to not generate gradient to allow implementation_selector # to work properly. # TODO(b/80444525): Remove the stop_gradient(). mask = array_ops . stop_gradient ( mask ) def input_right_padded ( ) : return cudnn_gru ( inputs = inputs , init_h = init_h , kernel = kernel , recurrent_kernel = recurrent_kernel , bias = bias , mask = mask , time_major = time_major , go_backwards = go_backwards ) def input_not_right_padded ( ) : return standard_gru ( inputs = inputs , init_h = init_h , kernel = kernel , recurrent_kernel = recurrent_kernel , bias = bias , mask = mask , time_major = time_major , go_backwards = go_backwards , activation = activation , recurrent_activation = recurrent_activation ) return control_flow_ops . cond ( is_sequence_right_padded ( mask , time_major ) , true_fn = input_right_padded , false_fn = input_not_right_padded ) # Each time a `tf.function` is called, we will give it a unique # identifiable API name, so that Grappler won't get confused when it # sees multiple GRU layers added into same graph, and it will be able # to pair up the different implementations across them. api_name = 'gru_' + str ( uuid . uuid4 ( ) ) defun_standard_gru = _generate_defun_backend ( api_name , _CPU_DEVICE_NAME , standard_gru ) defun_cudnn_gru = _generate_defun_backend ( api_name , _GPU_DEVICE_NAME , cudnn_gru_with_fallback ) # Call the normal GRU impl and register the CuDNN impl function. The # grappler will kick in during session execution to optimize the graph. last_output , outputs , new_h , runtime = defun_standard_gru ( * * params ) function . register ( defun_cudnn_gru , * * params ) return last_output , outputs , new_h , runtime