# Copyright 2019 The TensorFlow Authors. All Rights Reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ==============================================================================

"""Recurrent layers for TF 2.0.

"""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import uuid

from tensorflow.python.eager import context

from tensorflow.python.eager import function

from tensorflow.python.framework import constant_op

from tensorflow.python.framework import device

from tensorflow.python.framework import dtypes

from tensorflow.python.framework import ops

from tensorflow.python.keras import backend as K

from tensorflow.python.keras.engine.input_spec import InputSpec

from tensorflow.python.keras.layers import recurrent

from tensorflow.python.ops import array_ops

from tensorflow.python.ops import control_flow_ops

from tensorflow.python.ops import gen_cudnn_rnn_ops

from tensorflow.python.ops import math_ops

from tensorflow.python.ops import state_ops

from tensorflow.python.util.tf_export import keras_export

# The following string constants are used by Defun approach for unified backend

# of LSTM and GRU.

_DEFUN_API_NAME_ATTRIBUTE = 'api_implements'

_DEFUN_DEVICE_ATTRIBUTE = 'api_preferred_device'

_CPU_DEVICE_NAME = 'CPU'

_GPU_DEVICE_NAME = 'GPU'

# The following number constants are used to represent the runtime of the defun

# backend function. Since the CPU/GPU implementation are mathematically same, we

# need some signal for the function to indicate which function is executed. This

# is for testing purpose to verify the correctness of swapping backend function.

_RUNTIME_UNKNOWN = 0

_RUNTIME_CPU = 1

_RUNTIME_GPU = 2

@keras_export('keras.layers.GRUCell', v1=[])

class GRUCell(recurrent.GRUCell):

**kwargs)

@keras_export('keras.layers.GRU', v1=[])

class GRU(recurrent.DropoutRNNCellMixin, recurrent.GRU):

return last_output, outputs, runtime, states

def standard_gru(inputs, init_h, kernel, recurrent_kernel, bias, activation,

return last_output, outputs, new_states[0], _runtime(_RUNTIME_CPU)

def cudnn_gru(inputs, init_h, kernel, recurrent_kernel, bias, mask, time_major,

return last_output, outputs, h, _runtime(_RUNTIME_GPU)

def gru_with_backend_selection(

return last_output, outputs, new_h, runtime

@keras_export('keras.layers.LSTMCell', v1=[])

class LSTMCell(recurrent.LSTMCell):

**kwargs)

@keras_export('keras.layers.LSTM', v1=[])

class LSTM(recurrent.DropoutRNNCellMixin, recurrent.LSTM):

return output