# SPDX-License-Identifier: Apache-2.0


"""Unit Tests for custom rnns."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf

from tensorflow.python.ops import init_ops
from backend_test_base import Tf2OnnxBackendTestBase
from common import *  # pylint: disable=wildcard-import, unused-wildcard-import
from tf2onnx.tf_loader import is_tf2


# pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
# pylint: disable=abstract-method,arguments-differ

if is_tf2():
    BasicLSTMCell = tf.compat.v1.nn.rnn_cell.BasicLSTMCell
    LSTMCell = tf.compat.v1.nn.rnn_cell.LSTMCell
    GRUCell = tf.compat.v1.nn.rnn_cell.GRUCell
    RNNCell = tf.compat.v1.nn.rnn_cell.RNNCell
    MultiRNNCell = tf.compat.v1.nn.rnn_cell.MultiRNNCell
    dynamic_rnn = tf.compat.v1.nn.dynamic_rnn
    bidirectional_dynamic_rnn = tf.compat.v1.nn.bidirectional_dynamic_rnn
else:
    LSTMBlockCell = tf.contrib.rnn.LSTMBlockCell
    LSTMCell = tf.nn.rnn_cell.LSTMCell
    GRUCell = tf.nn.rnn_cell.LSTMCell
    RNNCell = tf.nn.rnn_cell.RNNCell
    MultiRNNCell = tf.contrib.rnn.MultiRNNCell
    dynamic_rnn = tf.nn.dynamic_rnn
    bidirectional_dynamic_rnn = tf.nn.bidirectional_dynamic_rnn


class CustomRnnCellTests(Tf2OnnxBackendTestBase):
    @check_opset_min_version(8, "Scan")
    @skip_tf2()
    def test_single_dynamic_custom_rnn(self):
        size = 5  # size of each model layer.
        batch_size = 1
        cell = GatedGRUCell(size)

        x_val = np.array([[1., 1.], [2., 2.], [3., 3.]], dtype=np.float32)
        x_val = np.stack([x_val] * batch_size)
        def func(x):
            xs, s = dynamic_rnn(cell=cell, dtype=tf.float32, inputs=x, time_major=False)
            return tf.identity(xs, name="output"), tf.identity(s, name="final_state")

        feed_dict = {"input_1:0": x_val}
        input_names_with_port = ["input_1:0"]
        output_names_with_port = ["output:0", "final_state:0"]
        self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port, 0.1)

    @check_opset_min_version(8, "Scan")
    @skip_tf2()
    def test_single_dynamic_custom_rnn_time_major(self):
        size = 5  # size of each model layer.
        batch_size = 1
        x_val = np.array([[1., 1.], [2., 2.], [3., 3.]], dtype=np.float32)
        x_val = np.stack([x_val] * batch_size)
        def func(x):
            cell = GatedGRUCell(size)
            xs, s = dynamic_rnn(cell=cell, dtype=tf.float32, inputs=x, time_major=True)
            return tf.identity(xs, name="output"), tf.identity(s, name="final_state")

        feed_dict = {"input_1:0": x_val}
        input_names_with_port = ["input_1:0"]
        output_names_with_port = ["output:0", "final_state:0"]
        self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port, 0.1)

    @check_opset_min_version(8, "Scan")
    @skip_tf2()
    def test_single_dynamic_custom_rnn_with_seq_length(self):
        units = 5
        batch_size = 6
        x_val = np.array([[1., 1.], [2., 2.], [3., 3.], [4., 4.], [5., 5.]], dtype=np.float32)
        x_val = np.stack([x_val] * batch_size)
        def func(x):
            # no scope
            cell = GatedGRUCell(units)
            outputs, cell_state = dynamic_rnn(
                cell,
                x,
                dtype=tf.float32,
                sequence_length=[4, 3, 4, 5, 2, 1])
            return tf.identity(outputs, name="output"), tf.identity(cell_state, name="cell_state")

        feed_dict = {"input_1:0": x_val}
        input_names_with_port = ["input_1:0"]
        output_names_with_port = ["output:0", "cell_state:0"]
        self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)

    @check_opset_min_version(8, "Scan")
    @skip_tf2()
    def test_single_dynamic_custom_rnn_with_non_const_seq_length(self):
        units = 5
        batch_size = 6
        x_val = np.array([[1., 1.], [2., 2.], [3., 3.], [4., 4.], [5., 5.]], dtype=np.float32)
        x_val = np.stack([x_val] * batch_size)
        y_val = np.array([4, 3, 4, 5, 2, 1], dtype=np.int32)
        def func(x, seq_length):
            # no scope
            cell = GatedGRUCell(units)
            outputs, cell_state = dynamic_rnn(
                cell,
                x,
                dtype=tf.float32,
                sequence_length=tf.identity(seq_length))
            return tf.identity(outputs, name="output"), tf.identity(cell_state, name="cell_state")

        feed_dict = {"input_1:0": x_val, "input_2:0": y_val}
        input_names_with_port = ["input_1:0", "input_2:0"]
        output_names_with_port = ["output:0", "cell_state:0"]
        self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)

    @check_opset_min_version(8, "Scan")
    @check_tf_min_version("1.8")
    @skip_tf2()
    def test_attention_wrapper_const_encoder(self):
        size = 5
        time_step = 3
        input_size = 4
        attn_size = size
        batch_size = 9
        # shape  [batch size, time step, size]
        # attention_state: usually the output of an RNN encoder.
        # This tensor should be shaped `[batch_size, max_time, ...]`.
        decoder_time_step = 6
        x_val = np.random.randn(decoder_time_step, input_size).astype('f')
        x_val = np.stack([x_val] * batch_size)

        attention_states = np.random.randn(batch_size, time_step, attn_size).astype('f')
        def func(x):
            attention_mechanism = tf.contrib.seq2seq.BahdanauAttention(attn_size, attention_states)

            match_input_fn = lambda curr_input, state: tf.concat([curr_input, state], axis=-1)
            cell = LSTMCell(size)
            match_cell_fw = tf.contrib.seq2seq.AttentionWrapper(cell,
                                                                attention_mechanism,
                                                                attention_layer_size=attn_size,
                                                                cell_input_fn=match_input_fn,
                                                                output_attention=False)
            output, attr_state = dynamic_rnn(match_cell_fw, x, dtype=tf.float32)

            return tf.identity(output, name="output"), tf.identity(attr_state.cell_state, name="final_state")

        feed_dict = {"input_1:0": x_val}
        input_names_with_port = ["input_1:0"]
        output_names_with_port = ["output:0"]
        output_names_with_port = ["output:0", "final_state:0"]
        self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port, 0.1)

    @check_opset_min_version(8, "Scan")
    @check_tf_min_version("1.8")
    @skip_tf2()
    def test_attention_wrapper_lstm_encoder(self):
        size = 5
        time_step = 3
        input_size = 4
        attn_size = size
        batch_size = 9

        # shape  [batch size, time step, size]
        # attention_state: usually the output of an RNN encoder.
        # This tensor should be shaped `[batch_size, max_time, ...]`
        encoder_time_step = time_step
        encoder_x_val = np.random.randn(encoder_time_step, input_size).astype('f')
        encoder_x_val = np.stack([encoder_x_val] * batch_size)
        decoder_time_step = 6
        decoder_x_val = np.random.randn(decoder_time_step, input_size).astype('f')
        decoder_x_val = np.stack([decoder_x_val] * batch_size)

        def func(encoder_x, decoder_x):
            encoder_cell = LSTMCell(size)
            output, attr_state = dynamic_rnn(encoder_cell, encoder_x, dtype=tf.float32)
            output_0 = tf.identity(output, name="output_0")
            attention_states = output
            attention_mechanism = tf.contrib.seq2seq.BahdanauAttention(attn_size,
                                                                       attention_states)

            match_input_fn = lambda curr_input, state: tf.concat([curr_input, state], axis=-1)
            cell = LSTMCell(size)
            match_cell_fw = tf.contrib.seq2seq.AttentionWrapper(cell,
                                                                attention_mechanism,
                                                                attention_layer_size=attn_size,
                                                                cell_input_fn=match_input_fn,
                                                                output_attention=False)

            output, attr_state = dynamic_rnn(match_cell_fw, decoder_x, dtype=tf.float32)

            return output_0, tf.identity(output, name="output"), tf.identity(attr_state.cell_state, name="final_state")

        feed_dict = {"input_1:0": encoder_x_val, "input_2:0": decoder_x_val}
        input_names_with_port = ["input_1:0", "input_2:0"]
        output_names_with_port = ["output_0:0", "output:0", "final_state:0"]
        self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port, 0.1)

    @check_opset_min_version(8, "Scan")
    @check_tf_min_version("1.8")
    @skip_tf2()
    def test_attention_wrapper_gru_encoder(self):
        size = 5
        time_step = 3
        input_size = 4
        attn_size = size
        batch_size = 9

        # shape  [batch size, time step, size]
        # attention_state: usually the output of an RNN encoder.
        # This tensor should be shaped `[batch_size, max_time, ...]`
        encoder_time_step = time_step
        encoder_x_val = np.random.randn(encoder_time_step, input_size).astype('f')
        encoder_x_val = np.stack([encoder_x_val] * batch_size)
        decoder_time_step = 6
        decoder_x_val = np.random.randn(decoder_time_step, input_size).astype('f')
        decoder_x_val = np.stack([decoder_x_val] * batch_size)

        def func(encoder_x, decoder_x):
            encoder_cell = GRUCell(size)
            output, attr_state = dynamic_rnn(encoder_cell, encoder_x, dtype=tf.float32)
            _ = tf.identity(output, name="output_0")
            attention_states = output
            attention_mechanism = tf.contrib.seq2seq.BahdanauAttention(attn_size,
                                                                       attention_states)

            match_input_fn = lambda curr_input, state: tf.concat([curr_input, state], axis=-1)
            cell = GRUCell(size)
            match_cell_fw = tf.contrib.seq2seq.AttentionWrapper(cell,
                                                                attention_mechanism,
                                                                attention_layer_size=attn_size,
                                                                cell_input_fn=match_input_fn,
                                                                output_attention=False)
            output, attr_state = dynamic_rnn(match_cell_fw, decoder_x, dtype=tf.float32)
            return tf.identity(output, name="output"), tf.identity(attr_state.cell_state, name="final_state")

        feed_dict = {"input_1:0": encoder_x_val, "input_2:0": decoder_x_val}
        input_names_with_port = ["input_1:0", "input_2:0"]
        output_names_with_port = ["output_0:0", "output:0", "final_state:0"]
        self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port, 0.1)

    @check_opset_min_version(8, "Scan")
    @check_tf_min_version("1.8")
    @skip_tf2()
    def test_attention_wrapper_lstm_encoder_input_has_none_dim(self):
        size = 5
        time_step = 3
        input_size = 4
        attn_size = size
        batch_size = 9

        # shape  [batch size, time step, size]
        # attention_state: usually the output of an RNN encoder.
        # This tensor should be shaped `[batch_size, max_time, ...]`
        encoder_time_step = time_step
        encoder_x_val = np.random.randn(encoder_time_step, input_size).astype('f')
        encoder_x_val = np.stack([encoder_x_val] * batch_size)
        decoder_time_step = 6
        decoder_x_val = np.random.randn(decoder_time_step, input_size).astype('f')
        decoder_x_val = np.stack([decoder_x_val] * batch_size)

        def func(encoder_x, decoder_x):
            encoder_cell = LSTMCell(size)
            output, attr_state = dynamic_rnn(encoder_cell, encoder_x, dtype=tf.float32)
            _ = tf.identity(output, name="output_0")
            attention_states = output
            attention_mechanism = tf.contrib.seq2seq.BahdanauAttention(attn_size,
                                                                       attention_states)

            match_input_fn = lambda curr_input, state: tf.concat([curr_input, state], axis=-1)
            cell = LSTMCell(size)
            match_cell_fw = tf.contrib.seq2seq.AttentionWrapper(cell,
                                                                attention_mechanism,
                                                                attention_layer_size=attn_size,
                                                                cell_input_fn=match_input_fn,
                                                                output_attention=False)

            output, attr_state = dynamic_rnn(match_cell_fw, decoder_x, dtype=tf.float32)

            return tf.identity(output, name="output"), tf.identity(attr_state.cell_state, name="final_state")

        feed_dict = {"input_1:0": encoder_x_val, "input_2:0": decoder_x_val}
        input_names_with_port = ["input_1:0", "input_2:0"]
        output_names_with_port = ["output_0:0", "output:0", "final_state:0"]

        self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port, 0.1)

    @check_opset_min_version(8, "Scan")
    @skip_tf2()
    def test_multi_rnn_lstm(self, state_is_tuple=True):
        units = 5
        batch_size = 6
        x_val = np.array([[1., 1.], [2., 2.], [3., 3.], [4., 4.]], dtype=np.float32)
        x_val = np.stack([x_val] * batch_size)
        def func(x):
            initializer = init_ops.constant_initializer(0.5)

            cell_0 = LSTMCell(units,
                              initializer=initializer,
                              state_is_tuple=state_is_tuple)

            cell_1 = LSTMCell(units,
                              initializer=initializer,
                              state_is_tuple=state_is_tuple)

            cell_2 = LSTMCell(units,
                              initializer=initializer,
                              state_is_tuple=state_is_tuple)

            cells = MultiRNNCell([cell_0, cell_1, cell_2], state_is_tuple=state_is_tuple)
            outputs, cell_state = dynamic_rnn(cells, x, dtype=tf.float32)
            return tf.identity(outputs, name="output"), tf.identity(cell_state, name="cell_state")

        input_names_with_port = ["input_1:0"]
        feed_dict = {"input_1:0": x_val}

        output_names_with_port = ["output:0", "cell_state:0"]
        self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port, rtol=1e-06)

    @check_opset_min_version(8, "Scan")
    @check_tf_min_version("1.8")
    @skip_opset(9, "ReverseSequence")
    @skip_tf2()
    @allow_missing_shapes("Missing RNN shape")
    def test_bidrectional_attention_wrapper_lstm_encoder(self):
        size = 30
        time_step = 3
        input_size = 4
        attn_size = size
        batch_size = 9

        # shape  [batch size, time step, size]
        # attention_state: usually the output of an RNN encoder.
        # This tensor should be shaped `[batch_size, max_time, ...]`
        encoder_time_step = time_step
        encoder_x_val = np.random.randn(encoder_time_step, input_size).astype('f')
        encoder_x_val = np.stack([encoder_x_val] * batch_size)
        decoder_time_step = 6
        decoder_x_val = np.random.randn(decoder_time_step, batch_size, input_size).astype('f')

        def func(encoder_x, decoder_x, seq_length):
            encoder_cell = LSTMCell(size)
            attention_states, _ = dynamic_rnn(encoder_cell, encoder_x, dtype=tf.float32)
            # [9, 3, 30], [9, 30]
            attention_mechanism = tf.contrib.seq2seq.BahdanauAttention(attn_size,
                                                                       attention_states)

            match_input_fn = lambda curr_input, state: tf.concat([curr_input, state], axis=-1)
            cell = LSTMCell(size)
            match_cell_fw = tf.contrib.seq2seq.AttentionWrapper(cell,
                                                                attention_mechanism,
                                                                attention_layer_size=attn_size,
                                                                cell_input_fn=match_input_fn,
                                                                output_attention=False)
            match_cell_bk = tf.contrib.seq2seq.AttentionWrapper(cell,
                                                                attention_mechanism,
                                                                attention_layer_size=attn_size,
                                                                cell_input_fn=match_input_fn,
                                                                output_attention=False)
            (match_output_fw, match_output_bk), (match_state_fw, match_state_bk) = \
                bidirectional_dynamic_rnn(cell_fw=match_cell_fw,
                                          cell_bw=match_cell_bk,
                                          inputs=decoder_x,
                                          sequence_length=tf.identity(seq_length),
                                          dtype=tf.float32,
                                          time_major=True)

            matched_output = tf.concat([match_output_fw, match_output_bk], axis=-1)
            matched_state = tf.concat([match_state_fw.cell_state, match_state_bk.cell_state], -1)
            return tf.identity(matched_output, name="output_0"), tf.identity(matched_state, name="final_state")

        feed_dict = {"input_1:0": encoder_x_val, "input_2:0": decoder_x_val,
                     "input_3:0": np.array([6, 5, 4, 3, 2, 1, 2, 3, 6], dtype=np.int32)}
        input_names_with_port = ["input_1:0", "input_2:0", "input_3:0"]
        output_names_with_port = ["output_0:0", "final_state:0"]
        self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port, 0.1)


class GatedGRUCell(RNNCell):
    def __init__(self, hidden_dim, reuse=None):
        super().__init__(self, _reuse=reuse)
        self._num_units = hidden_dim
        self._activation = tf.tanh

    @property
    def state_size(self):
        return self._num_units

    @property
    def output_size(self):
        return self._num_units

    def call(self, inputs, state):
        # inputs shape: [batch size, time step, input size] = [1, 3, 2]
        # num_units: 5
        # W shape: [2, 3 * 5] = [2, 15]
        # U shape: [5, 3 * 5] = [5, 15]
        # b shape: [1, 3 * 5] = [1, 15]
        # state shape: [batch size, state size] = [1, 5]

        input_dim = inputs.get_shape()[-1]
        assert input_dim is not None, "input dimension must be defined"
        # W = tf.get_variable(name="W", shape=[input_dim, 3 * self._num_units], dtype=tf.float32)
        W = np.arange(30.0, dtype=np.float32).reshape((2, 15))
        # U = tf.get_variable(name='U', shape=[self._num_units, 3 * self._num_units], dtype=tf.float32)
        U = np.arange(75.0, dtype=np.float32).reshape((5, 15))
        # b = tf.get_variable(name='b', shape=[1, 3 * self._num_units], dtype=tf.float32)
        b = np.arange(15.0, dtype=np.float32).reshape((1, 15))

        xw = tf.split(tf.matmul(inputs, W) + b, 3, 1)
        hu = tf.split(tf.matmul(state, U), 3, 1)
        r = tf.sigmoid(xw[0] + hu[0])
        z = tf.sigmoid(xw[1] + hu[1])
        h1 = self._activation(xw[2] + r * hu[2])
        next_h = h1 * (1 - z) + state * z
        return next_h, next_h


if __name__ == '__main__':
    unittest_main()

[((22, 3, 22, 11), 'tf2onnx.tf_loader.is_tf2', 'is_tf2', ({}, {}), '()', False, 'from tf2onnx.tf_loader import is_tf2\n'), ((48, 16, 48, 74), 'numpy.array', 'np.array', (), '', True, 'import numpy as np\n'), ((49, 16, 49, 46), 'numpy.stack', 'np.stack', ({(49, 25, 49, 45): '[x_val] * batch_size'}, {}), '([x_val] * batch_size)', True, 'import numpy as np\n'), ((64, 16, 64, 74), 'numpy.array', 'np.array', (), '', True, 'import numpy as np\n'), ((65, 16, 65, 46), 'numpy.stack', 'np.stack', ({(65, 25, 65, 45): '[x_val] * batch_size'}, {}), '([x_val] * batch_size)', True, 'import numpy as np\n'), ((81, 16, 81, 94), 'numpy.array', 'np.array', (), '', True, 'import numpy as np\n'), ((82, 16, 82, 46), 'numpy.stack', 'np.stack', ({(82, 25, 82, 45): '[x_val] * batch_size'}, {}), '([x_val] * batch_size)', True, 'import numpy as np\n'), ((103, 16, 103, 94), 'numpy.array', 'np.array', (), '', True, 'import numpy as np\n'), ((104, 16, 104, 46), 'numpy.stack', 'np.stack', ({(104, 25, 104, 45): '[x_val] * batch_size'}, {}), '([x_val] * batch_size)', True, 'import numpy as np\n'), ((105, 16, 105, 60), 'numpy.array', 'np.array', (), '', True, 'import numpy as np\n'), ((135, 16, 135, 46), 'numpy.stack', 'np.stack', ({(135, 25, 135, 45): '[x_val] * batch_size'}, {}), '([x_val] * batch_size)', True, 'import numpy as np\n'), ((173, 24, 173, 62), 'numpy.stack', 'np.stack', ({(173, 33, 173, 61): '[encoder_x_val] * batch_size'}, {}), '([encoder_x_val] * batch_size)', True, 'import numpy as np\n'), ((176, 24, 176, 62), 'numpy.stack', 'np.stack', ({(176, 33, 176, 61): '[decoder_x_val] * batch_size'}, {}), '([decoder_x_val] * batch_size)', True, 'import numpy as np\n'), ((218, 24, 218, 62), 'numpy.stack', 'np.stack', ({(218, 33, 218, 61): '[encoder_x_val] * batch_size'}, {}), '([encoder_x_val] * batch_size)', True, 'import numpy as np\n'), ((221, 24, 221, 62), 'numpy.stack', 'np.stack', ({(221, 33, 221, 61): '[decoder_x_val] * batch_size'}, {}), '([decoder_x_val] * batch_size)', True, 'import numpy as np\n'), ((261, 24, 261, 62), 'numpy.stack', 'np.stack', ({(261, 33, 261, 61): '[encoder_x_val] * batch_size'}, {}), '([encoder_x_val] * batch_size)', True, 'import numpy as np\n'), ((264, 24, 264, 62), 'numpy.stack', 'np.stack', ({(264, 33, 264, 61): '[decoder_x_val] * batch_size'}, {}), '([decoder_x_val] * batch_size)', True, 'import numpy as np\n'), ((297, 16, 297, 84), 'numpy.array', 'np.array', (), '', True, 'import numpy as np\n'), ((298, 16, 298, 46), 'numpy.stack', 'np.stack', ({(298, 25, 298, 45): '[x_val] * batch_size'}, {}), '([x_val] * batch_size)', True, 'import numpy as np\n'), ((341, 24, 341, 62), 'numpy.stack', 'np.stack', ({(341, 33, 341, 61): '[encoder_x_val] * batch_size'}, {}), '([encoder_x_val] * batch_size)', True, 'import numpy as np\n'), ((416, 12, 416, 37), 'tensorflow.sigmoid', 'tf.sigmoid', ({(416, 23, 416, 36): 'xw[0] + hu[0]'}, {}), '(xw[0] + hu[0])', True, 'import tensorflow as tf\n'), ((417, 12, 417, 37), 'tensorflow.sigmoid', 'tf.sigmoid', ({(417, 23, 417, 36): 'xw[1] + hu[1]'}, {}), '(xw[1] + hu[1])', True, 'import tensorflow as tf\n'), ((139, 34, 139, 99), 'tensorflow.contrib.seq2seq.BahdanauAttention', 'tf.contrib.seq2seq.BahdanauAttention', ({(139, 71, 139, 80): 'attn_size', (139, 82, 139, 98): 'attention_states'}, {}), '(attn_size, attention_states)', True, 'import tensorflow as tf\n'), ((143, 28, 147, 87), 'tensorflow.contrib.seq2seq.AttentionWrapper', 'tf.contrib.seq2seq.AttentionWrapper', (), '', True, 'import tensorflow as tf\n'), ((181, 23, 181, 59), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((183, 34, 184, 88), 'tensorflow.contrib.seq2seq.BahdanauAttention', 'tf.contrib.seq2seq.BahdanauAttention', ({(183, 71, 183, 80): 'attn_size', (184, 71, 184, 87): 'attention_states'}, {}), '(attn_size, attention_states)', True, 'import tensorflow as tf\n'), ((188, 28, 192, 87), 'tensorflow.contrib.seq2seq.AttentionWrapper', 'tf.contrib.seq2seq.AttentionWrapper', (), '', True, 'import tensorflow as tf\n'), ((226, 16, 226, 52), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((228, 34, 229, 88), 'tensorflow.contrib.seq2seq.BahdanauAttention', 'tf.contrib.seq2seq.BahdanauAttention', ({(228, 71, 228, 80): 'attn_size', (229, 71, 229, 87): 'attention_states'}, {}), '(attn_size, attention_states)', True, 'import tensorflow as tf\n'), ((233, 28, 237, 87), 'tensorflow.contrib.seq2seq.AttentionWrapper', 'tf.contrib.seq2seq.AttentionWrapper', (), '', True, 'import tensorflow as tf\n'), ((269, 16, 269, 52), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((271, 34, 272, 88), 'tensorflow.contrib.seq2seq.BahdanauAttention', 'tf.contrib.seq2seq.BahdanauAttention', ({(271, 71, 271, 80): 'attn_size', (272, 71, 272, 87): 'attention_states'}, {}), '(attn_size, attention_states)', True, 'import tensorflow as tf\n'), ((276, 28, 280, 87), 'tensorflow.contrib.seq2seq.AttentionWrapper', 'tf.contrib.seq2seq.AttentionWrapper', (), '', True, 'import tensorflow as tf\n'), ((300, 26, 300, 60), 'tensorflow.python.ops.init_ops.constant_initializer', 'init_ops.constant_initializer', ({(300, 56, 300, 59): '0.5'}, {}), '(0.5)', False, 'from tensorflow.python.ops import init_ops\n'), ((349, 34, 350, 88), 'tensorflow.contrib.seq2seq.BahdanauAttention', 'tf.contrib.seq2seq.BahdanauAttention', ({(349, 71, 349, 80): 'attn_size', (350, 71, 350, 87): 'attention_states'}, {}), '(attn_size, attention_states)', True, 'import tensorflow as tf\n'), ((354, 28, 358, 87), 'tensorflow.contrib.seq2seq.AttentionWrapper', 'tf.contrib.seq2seq.AttentionWrapper', (), '', True, 'import tensorflow as tf\n'), ((359, 28, 363, 87), 'tensorflow.contrib.seq2seq.AttentionWrapper', 'tf.contrib.seq2seq.AttentionWrapper', (), '', True, 'import tensorflow as tf\n'), ((372, 29, 372, 83), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((373, 28, 373, 97), 'tensorflow.concat', 'tf.concat', ({(373, 38, 373, 92): '[match_state_fw.cell_state, match_state_bk.cell_state]', (373, 94, 373, 96): '-1'}, {}), '([match_state_fw.cell_state, match_state_bk.cell_state], -1)', True, 'import tensorflow as tf\n'), ((377, 34, 377, 87), 'numpy.array', 'np.array', (), '', True, 'import numpy as np\n'), ((415, 22, 415, 41), 'tensorflow.matmul', 'tf.matmul', ({(415, 32, 415, 37): 'state', (415, 39, 415, 40): 'U'}, {}), '(state, U)', True, 'import tensorflow as tf\n'), ((52, 19, 52, 49), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((52, 51, 52, 85), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((69, 19, 69, 49), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((69, 51, 69, 85), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((91, 19, 91, 54), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((91, 56, 91, 98), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((114, 19, 114, 54), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((114, 56, 114, 98), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((134, 16, 134, 62), 'numpy.random.randn', 'np.random.randn', ({(134, 32, 134, 49): 'decoder_time_step', (134, 51, 134, 61): 'input_size'}, {}), '(decoder_time_step, input_size)', True, 'import numpy as np\n'), ((137, 27, 137, 76), 'numpy.random.randn', 'np.random.randn', ({(137, 43, 137, 53): 'batch_size', (137, 55, 137, 64): 'time_step', (137, 66, 137, 75): 'attn_size'}, {}), '(batch_size, time_step, attn_size)', True, 'import numpy as np\n'), ((141, 55, 141, 94), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((150, 19, 150, 53), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((150, 55, 150, 109), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((172, 24, 172, 70), 'numpy.random.randn', 'np.random.randn', ({(172, 40, 172, 57): 'encoder_time_step', (172, 59, 172, 69): 'input_size'}, {}), '(encoder_time_step, input_size)', True, 'import numpy as np\n'), ((175, 24, 175, 70), 'numpy.random.randn', 'np.random.randn', ({(175, 40, 175, 57): 'decoder_time_step', (175, 59, 175, 69): 'input_size'}, {}), '(decoder_time_step, input_size)', True, 'import numpy as np\n'), ((186, 55, 186, 94), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((196, 29, 196, 63), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((196, 65, 196, 119), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((217, 24, 217, 70), 'numpy.random.randn', 'np.random.randn', ({(217, 40, 217, 57): 'encoder_time_step', (217, 59, 217, 69): 'input_size'}, {}), '(encoder_time_step, input_size)', True, 'import numpy as np\n'), ((220, 24, 220, 70), 'numpy.random.randn', 'np.random.randn', ({(220, 40, 220, 57): 'decoder_time_step', (220, 59, 220, 69): 'input_size'}, {}), '(decoder_time_step, input_size)', True, 'import numpy as np\n'), ((231, 55, 231, 94), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((239, 19, 239, 53), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((239, 55, 239, 109), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((260, 24, 260, 70), 'numpy.random.randn', 'np.random.randn', ({(260, 40, 260, 57): 'encoder_time_step', (260, 59, 260, 69): 'input_size'}, {}), '(encoder_time_step, input_size)', True, 'import numpy as np\n'), ((263, 24, 263, 70), 'numpy.random.randn', 'np.random.randn', ({(263, 40, 263, 57): 'decoder_time_step', (263, 59, 263, 69): 'input_size'}, {}), '(decoder_time_step, input_size)', True, 'import numpy as np\n'), ((274, 55, 274, 94), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((284, 19, 284, 53), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((284, 55, 284, 109), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((316, 19, 316, 54), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((316, 56, 316, 98), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((340, 24, 340, 70), 'numpy.random.randn', 'np.random.randn', ({(340, 40, 340, 57): 'encoder_time_step', (340, 59, 340, 69): 'input_size'}, {}), '(encoder_time_step, input_size)', True, 'import numpy as np\n'), ((343, 24, 343, 82), 'numpy.random.randn', 'np.random.randn', ({(343, 40, 343, 57): 'decoder_time_step', (343, 59, 343, 69): 'batch_size', (343, 71, 343, 81): 'input_size'}, {}), '(decoder_time_step, batch_size, input_size)', True, 'import numpy as np\n'), ((352, 55, 352, 94), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((374, 19, 374, 63), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((374, 65, 374, 111), 'tensorflow.identity', 'tf.identity', (), '', True, 'import tensorflow as tf\n'), ((408, 12, 408, 45), 'numpy.arange', 'np.arange', (), '', True, 'import numpy as np\n'), ((410, 12, 410, 45), 'numpy.arange', 'np.arange', (), '', True, 'import numpy as np\n'), ((412, 12, 412, 45), 'numpy.arange', 'np.arange', (), '', True, 'import numpy as np\n'), ((414, 22, 414, 42), 'tensorflow.matmul', 'tf.matmul', ({(414, 32, 414, 38): 'inputs', (414, 40, 414, 41): 'W'}, {}), '(inputs, W)', True, 'import tensorflow as tf\n'), ((113, 32, 113, 55), 'tensorflow.identity', 'tf.identity', ({(113, 44, 113, 54): 'seq_length'}, {}), '(seq_length)', True, 'import tensorflow as tf\n'), ((368, 58, 368, 81), 'tensorflow.identity', 'tf.identity', ({(368, 70, 368, 80): 'seq_length'}, {}), '(seq_length)', True, 'import tensorflow as tf\n')]