# Copyright 2019 The Magenta Authors.
#
# 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.

"""SketchRNN RNN definition."""

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

import numpy as np
import tensorflow.compat.v1 as tf
from tensorflow.contrib import rnn as contrib_rnn


def orthogonal(shape):
  """Orthogonal initilaizer."""
  flat_shape = (shape[0], np.prod(shape[1:]))
  a = np.random.normal(0.0, 1.0, flat_shape)
  u, _, v = np.linalg.svd(a, full_matrices=False)
  q = u if u.shape == flat_shape else v
  return q.reshape(shape)


def orthogonal_initializer(scale=1.0):
  """Orthogonal initializer."""
  def _initializer(shape, dtype=tf.float32,
                   partition_info=None):  # pylint: disable=unused-argument
    return tf.constant(orthogonal(shape) * scale, dtype)

  return _initializer


def lstm_ortho_initializer(scale=1.0):
  """LSTM orthogonal initializer."""
  def _initializer(shape, dtype=tf.float32,
                   partition_info=None):  # pylint: disable=unused-argument
    size_x = shape[0]
    size_h = shape[1] // 4  # assumes lstm.
    t = np.zeros(shape)
    t[:, :size_h] = orthogonal([size_x, size_h]) * scale
    t[:, size_h:size_h * 2] = orthogonal([size_x, size_h]) * scale
    t[:, size_h * 2:size_h * 3] = orthogonal([size_x, size_h]) * scale
    t[:, size_h * 3:] = orthogonal([size_x, size_h]) * scale
    return tf.constant(t, dtype)

  return _initializer


class LSTMCell(contrib_rnn.RNNCell):
  """Vanilla LSTM cell.

  Uses ortho initializer, and also recurrent dropout without memory loss
  (https://arxiv.org/abs/1603.05118)
  """

  def __init__(self,
               num_units,
               forget_bias=1.0,
               use_recurrent_dropout=False,
               dropout_keep_prob=0.9):
    self.num_units = num_units
    self.forget_bias = forget_bias
    self.use_recurrent_dropout = use_recurrent_dropout
    self.dropout_keep_prob = dropout_keep_prob

  @property
  def state_size(self):
    return 2 * self.num_units

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

  def get_output(self, state):
    unused_c, h = tf.split(state, 2, 1)
    return h

  def __call__(self, x, state, scope=None):
    with tf.variable_scope(scope or type(self).__name__):
      c, h = tf.split(state, 2, 1)

      x_size = x.get_shape().as_list()[1]

      w_init = None  # uniform

      h_init = lstm_ortho_initializer(1.0)

      # Keep W_xh and W_hh separate here as well to use different init methods.
      w_xh = tf.get_variable(
          'W_xh', [x_size, 4 * self.num_units], initializer=w_init)
      w_hh = tf.get_variable(
          'W_hh', [self.num_units, 4 * self.num_units], initializer=h_init)
      bias = tf.get_variable(
          'bias', [4 * self.num_units],
          initializer=tf.constant_initializer(0.0))

      concat = tf.concat([x, h], 1)
      w_full = tf.concat([w_xh, w_hh], 0)
      hidden = tf.matmul(concat, w_full) + bias

      i, j, f, o = tf.split(hidden, 4, 1)

      if self.use_recurrent_dropout:
        g = tf.nn.dropout(tf.tanh(j), self.dropout_keep_prob)
      else:
        g = tf.tanh(j)

      new_c = c * tf.sigmoid(f + self.forget_bias) + tf.sigmoid(i) * g
      new_h = tf.tanh(new_c) * tf.sigmoid(o)

      return new_h, tf.concat([new_c, new_h], 1)  # fuk tuples.


def layer_norm_all(h,
                   batch_size,
                   base,
                   num_units,
                   scope='layer_norm',
                   reuse=False,
                   gamma_start=1.0,
                   epsilon=1e-3,
                   use_bias=True):
  """Layer Norm (faster version, but not using defun)."""
  # Performs layer norm on multiple base at once (ie, i, g, j, o for lstm)
  # Reshapes h in to perform layer norm in parallel
  h_reshape = tf.reshape(h, [batch_size, base, num_units])
  mean = tf.reduce_mean(h_reshape, [2], keep_dims=True)
  var = tf.reduce_mean(tf.square(h_reshape - mean), [2], keep_dims=True)
  epsilon = tf.constant(epsilon)
  rstd = tf.rsqrt(var + epsilon)
  h_reshape = (h_reshape - mean) * rstd
  # reshape back to original
  h = tf.reshape(h_reshape, [batch_size, base * num_units])
  with tf.variable_scope(scope):
    if reuse:
      tf.get_variable_scope().reuse_variables()
    gamma = tf.get_variable(
        'ln_gamma', [4 * num_units],
        initializer=tf.constant_initializer(gamma_start))
    if use_bias:
      beta = tf.get_variable(
          'ln_beta', [4 * num_units], initializer=tf.constant_initializer(0.0))
  if use_bias:
    return gamma * h + beta
  return gamma * h


def layer_norm(x,
               num_units,
               scope='layer_norm',
               reuse=False,
               gamma_start=1.0,
               epsilon=1e-3,
               use_bias=True):
  """Calculate layer norm."""
  axes = [1]
  mean = tf.reduce_mean(x, axes, keep_dims=True)
  x_shifted = x - mean
  var = tf.reduce_mean(tf.square(x_shifted), axes, keep_dims=True)
  inv_std = tf.rsqrt(var + epsilon)
  with tf.variable_scope(scope):
    if reuse:
      tf.get_variable_scope().reuse_variables()
    gamma = tf.get_variable(
        'ln_gamma', [num_units],
        initializer=tf.constant_initializer(gamma_start))
    if use_bias:
      beta = tf.get_variable(
          'ln_beta', [num_units], initializer=tf.constant_initializer(0.0))
  output = gamma * (x_shifted) * inv_std
  if use_bias:
    output += beta
  return output


def raw_layer_norm(x, epsilon=1e-3):
  axes = [1]
  mean = tf.reduce_mean(x, axes, keep_dims=True)
  std = tf.sqrt(
      tf.reduce_mean(tf.square(x - mean), axes, keep_dims=True) + epsilon)
  output = (x - mean) / (std)
  return output


def super_linear(x,
                 output_size,
                 scope=None,
                 reuse=False,
                 init_w='ortho',
                 weight_start=0.0,
                 use_bias=True,
                 bias_start=0.0,
                 input_size=None):
  """Performs linear operation. Uses ortho init defined earlier."""
  shape = x.get_shape().as_list()
  with tf.variable_scope(scope or 'linear'):
    if reuse:
      tf.get_variable_scope().reuse_variables()

    w_init = None  # uniform
    if input_size is None:
      x_size = shape[1]
    else:
      x_size = input_size
    if init_w == 'zeros':
      w_init = tf.constant_initializer(0.0)
    elif init_w == 'constant':
      w_init = tf.constant_initializer(weight_start)
    elif init_w == 'gaussian':
      w_init = tf.random_normal_initializer(stddev=weight_start)
    elif init_w == 'ortho':
      w_init = lstm_ortho_initializer(1.0)

    w = tf.get_variable(
        'super_linear_w', [x_size, output_size], tf.float32, initializer=w_init)
    if use_bias:
      b = tf.get_variable(
          'super_linear_b', [output_size],
          tf.float32,
          initializer=tf.constant_initializer(bias_start))
      return tf.matmul(x, w) + b
    return tf.matmul(x, w)


class LayerNormLSTMCell(contrib_rnn.RNNCell):
  """Layer-Norm, with Ortho Init. and Recurrent Dropout without Memory Loss.

  https://arxiv.org/abs/1607.06450 - Layer Norm
  https://arxiv.org/abs/1603.05118 - Recurrent Dropout without Memory Loss
  """

  def __init__(self,
               num_units,
               forget_bias=1.0,
               use_recurrent_dropout=False,
               dropout_keep_prob=0.90):
    """Initialize the Layer Norm LSTM cell.

    Args:
      num_units: int, The number of units in the LSTM cell.
      forget_bias: float, The bias added to forget gates (default 1.0).
      use_recurrent_dropout: Whether to use Recurrent Dropout (default False)
      dropout_keep_prob: float, dropout keep probability (default 0.90)
    """
    self.num_units = num_units
    self.forget_bias = forget_bias
    self.use_recurrent_dropout = use_recurrent_dropout
    self.dropout_keep_prob = dropout_keep_prob

  @property
  def input_size(self):
    return self.num_units

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

  @property
  def state_size(self):
    return 2 * self.num_units

  def get_output(self, state):
    h, unused_c = tf.split(state, 2, 1)
    return h

  def __call__(self, x, state, timestep=0, scope=None):
    with tf.variable_scope(scope or type(self).__name__):
      h, c = tf.split(state, 2, 1)

      h_size = self.num_units
      x_size = x.get_shape().as_list()[1]
      batch_size = x.get_shape().as_list()[0]

      w_init = None  # uniform

      h_init = lstm_ortho_initializer(1.0)

      w_xh = tf.get_variable(
          'W_xh', [x_size, 4 * self.num_units], initializer=w_init)
      w_hh = tf.get_variable(
          'W_hh', [self.num_units, 4 * self.num_units], initializer=h_init)

      concat = tf.concat([x, h], 1)  # concat for speed.
      w_full = tf.concat([w_xh, w_hh], 0)
      concat = tf.matmul(concat, w_full)  #+ bias # live life without garbage.

      # i = input_gate, j = new_input, f = forget_gate, o = output_gate
      concat = layer_norm_all(concat, batch_size, 4, h_size, 'ln_all')
      i, j, f, o = tf.split(concat, 4, 1)

      if self.use_recurrent_dropout:
        g = tf.nn.dropout(tf.tanh(j), self.dropout_keep_prob)
      else:
        g = tf.tanh(j)

      new_c = c * tf.sigmoid(f + self.forget_bias) + tf.sigmoid(i) * g
      new_h = tf.tanh(layer_norm(new_c, h_size, 'ln_c')) * tf.sigmoid(o)

    return new_h, tf.concat([new_h, new_c], 1)


class HyperLSTMCell(contrib_rnn.RNNCell):
  """HyperLSTM with Ortho Init, Layer Norm, Recurrent Dropout, no Memory Loss.

  https://arxiv.org/abs/1609.09106
  http://blog.otoro.net/2016/09/28/hyper-networks/
  """

  def __init__(self,
               num_units,
               forget_bias=1.0,
               use_recurrent_dropout=False,
               dropout_keep_prob=0.90,
               use_layer_norm=True,
               hyper_num_units=256,
               hyper_embedding_size=32,
               hyper_use_recurrent_dropout=False):
    """Initialize the Layer Norm HyperLSTM cell.

    Args:
      num_units: int, The number of units in the LSTM cell.
      forget_bias: float, The bias added to forget gates (default 1.0).
      use_recurrent_dropout: Whether to use Recurrent Dropout (default False)
      dropout_keep_prob: float, dropout keep probability (default 0.90)
      use_layer_norm: boolean. (default True)
        Controls whether we use LayerNorm layers in main LSTM & HyperLSTM cell.
      hyper_num_units: int, number of units in HyperLSTM cell.
        (default is 128, recommend experimenting with 256 for larger tasks)
      hyper_embedding_size: int, size of signals emitted from HyperLSTM cell.
        (default is 16, recommend trying larger values for large datasets)
      hyper_use_recurrent_dropout: boolean. (default False)
        Controls whether HyperLSTM cell also uses recurrent dropout.
        Recommend turning this on only if hyper_num_units becomes large (>= 512)
    """
    self.num_units = num_units
    self.forget_bias = forget_bias
    self.use_recurrent_dropout = use_recurrent_dropout
    self.dropout_keep_prob = dropout_keep_prob
    self.use_layer_norm = use_layer_norm
    self.hyper_num_units = hyper_num_units
    self.hyper_embedding_size = hyper_embedding_size
    self.hyper_use_recurrent_dropout = hyper_use_recurrent_dropout

    self.total_num_units = self.num_units + self.hyper_num_units

    if self.use_layer_norm:
      cell_fn = LayerNormLSTMCell
    else:
      cell_fn = LSTMCell
    self.hyper_cell = cell_fn(
        hyper_num_units,
        use_recurrent_dropout=hyper_use_recurrent_dropout,
        dropout_keep_prob=dropout_keep_prob)

  @property
  def input_size(self):
    return self._input_size

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

  @property
  def state_size(self):
    return 2 * self.total_num_units

  def get_output(self, state):
    total_h, unused_total_c = tf.split(state, 2, 1)
    h = total_h[:, 0:self.num_units]
    return h

  def hyper_norm(self, layer, scope='hyper', use_bias=True):
    num_units = self.num_units
    embedding_size = self.hyper_embedding_size
    # recurrent batch norm init trick (https://arxiv.org/abs/1603.09025).
    init_gamma = 0.10  # cooijmans' da man.
    with tf.variable_scope(scope):
      zw = super_linear(
          self.hyper_output,
          embedding_size,
          init_w='constant',
          weight_start=0.00,
          use_bias=True,
          bias_start=1.0,
          scope='zw')
      alpha = super_linear(
          zw,
          num_units,
          init_w='constant',
          weight_start=init_gamma / embedding_size,
          use_bias=False,
          scope='alpha')
      result = tf.multiply(alpha, layer)
      if use_bias:
        zb = super_linear(
            self.hyper_output,
            embedding_size,
            init_w='gaussian',
            weight_start=0.01,
            use_bias=False,
            bias_start=0.0,
            scope='zb')
        beta = super_linear(
            zb,
            num_units,
            init_w='constant',
            weight_start=0.00,
            use_bias=False,
            scope='beta')
        result += beta
    return result

  def __call__(self, x, state, timestep=0, scope=None):
    with tf.variable_scope(scope or type(self).__name__):
      total_h, total_c = tf.split(state, 2, 1)
      h = total_h[:, 0:self.num_units]
      c = total_c[:, 0:self.num_units]
      self.hyper_state = tf.concat(
          [total_h[:, self.num_units:], total_c[:, self.num_units:]], 1)

      batch_size = x.get_shape().as_list()[0]
      x_size = x.get_shape().as_list()[1]
      self._input_size = x_size

      w_init = None  # uniform

      h_init = lstm_ortho_initializer(1.0)

      w_xh = tf.get_variable(
          'W_xh', [x_size, 4 * self.num_units], initializer=w_init)
      w_hh = tf.get_variable(
          'W_hh', [self.num_units, 4 * self.num_units], initializer=h_init)
      bias = tf.get_variable(
          'bias', [4 * self.num_units],
          initializer=tf.constant_initializer(0.0))

      # concatenate the input and hidden states for hyperlstm input
      hyper_input = tf.concat([x, h], 1)
      hyper_output, hyper_new_state = self.hyper_cell(hyper_input,
                                                      self.hyper_state)
      self.hyper_output = hyper_output
      self.hyper_state = hyper_new_state

      xh = tf.matmul(x, w_xh)
      hh = tf.matmul(h, w_hh)

      # split Wxh contributions
      ix, jx, fx, ox = tf.split(xh, 4, 1)
      ix = self.hyper_norm(ix, 'hyper_ix', use_bias=False)
      jx = self.hyper_norm(jx, 'hyper_jx', use_bias=False)
      fx = self.hyper_norm(fx, 'hyper_fx', use_bias=False)
      ox = self.hyper_norm(ox, 'hyper_ox', use_bias=False)

      # split Whh contributions
      ih, jh, fh, oh = tf.split(hh, 4, 1)
      ih = self.hyper_norm(ih, 'hyper_ih', use_bias=True)
      jh = self.hyper_norm(jh, 'hyper_jh', use_bias=True)
      fh = self.hyper_norm(fh, 'hyper_fh', use_bias=True)
      oh = self.hyper_norm(oh, 'hyper_oh', use_bias=True)

      # split bias
      ib, jb, fb, ob = tf.split(bias, 4, 0)  # bias is to be broadcasted.

      # i = input_gate, j = new_input, f = forget_gate, o = output_gate
      i = ix + ih + ib
      j = jx + jh + jb
      f = fx + fh + fb
      o = ox + oh + ob

      if self.use_layer_norm:
        concat = tf.concat([i, j, f, o], 1)
        concat = layer_norm_all(concat, batch_size, 4, self.num_units, 'ln_all')
        i, j, f, o = tf.split(concat, 4, 1)

      if self.use_recurrent_dropout:
        g = tf.nn.dropout(tf.tanh(j), self.dropout_keep_prob)
      else:
        g = tf.tanh(j)

      new_c = c * tf.sigmoid(f + self.forget_bias) + tf.sigmoid(i) * g
      new_h = tf.tanh(layer_norm(new_c, self.num_units, 'ln_c')) * tf.sigmoid(o)

      hyper_h, hyper_c = tf.split(hyper_new_state, 2, 1)
      new_total_h = tf.concat([new_h, hyper_h], 1)
      new_total_c = tf.concat([new_c, hyper_c], 1)
      new_total_state = tf.concat([new_total_h, new_total_c], 1)
    return new_h, new_total_state

[((29, 6, 29, 44), 'numpy.random.normal', 'np.random.normal', ({(29, 23, 29, 26): '0.0', (29, 28, 29, 31): '1.0', (29, 33, 29, 43): 'flat_shape'}, {}), '(0.0, 1.0, flat_shape)', True, 'import numpy as np\n'), ((30, 12, 30, 49), 'numpy.linalg.svd', 'np.linalg.svd', (), '', True, 'import numpy as np\n'), ((137, 14, 137, 58), 'tensorflow.compat.v1.reshape', 'tf.reshape', ({(137, 25, 137, 26): 'h', (137, 28, 137, 57): '[batch_size, base, num_units]'}, {}), '(h, [batch_size, base, num_units])', True, 'import tensorflow.compat.v1 as tf\n'), ((138, 9, 138, 55), 'tensorflow.compat.v1.reduce_mean', 'tf.reduce_mean', (), '', True, 'import tensorflow.compat.v1 as tf\n'), ((140, 12, 140, 32), 'tensorflow.compat.v1.constant', 'tf.constant', ({(140, 24, 140, 31): 'epsilon'}, {}), '(epsilon)', True, 'import tensorflow.compat.v1 as tf\n'), ((141, 9, 141, 32), 'tensorflow.compat.v1.rsqrt', 'tf.rsqrt', ({(141, 18, 141, 31): 'var + epsilon'}, {}), '(var + epsilon)', True, 'import tensorflow.compat.v1 as tf\n'), ((144, 6, 144, 59), 'tensorflow.compat.v1.reshape', 'tf.reshape', ({(144, 17, 144, 26): 'h_reshape', (144, 28, 144, 58): '[batch_size, base * num_units]'}, {}), '(h_reshape, [batch_size, base * num_units])', True, 'import tensorflow.compat.v1 as tf\n'), ((168, 9, 168, 48), 'tensorflow.compat.v1.reduce_mean', 'tf.reduce_mean', (), '', True, 'import tensorflow.compat.v1 as tf\n'), ((171, 12, 171, 35), 'tensorflow.compat.v1.rsqrt', 'tf.rsqrt', ({(171, 21, 171, 34): 'var + epsilon'}, {}), '(var + epsilon)', True, 'import tensorflow.compat.v1 as tf\n'), ((189, 9, 189, 48), 'tensorflow.compat.v1.reduce_mean', 'tf.reduce_mean', (), '', True, 'import tensorflow.compat.v1 as tf\n'), ((28, 26, 28, 44), 'numpy.prod', 'np.prod', ({(28, 34, 28, 43): 'shape[1:]'}, {}), '(shape[1:])', True, 'import numpy as np\n'), ((50, 8, 50, 23), 'numpy.zeros', 'np.zeros', ({(50, 17, 50, 22): 'shape'}, {}), '(shape)', True, 'import numpy as np\n'), ((55, 11, 55, 32), 'tensorflow.compat.v1.constant', 'tf.constant', ({(55, 23, 55, 24): 't', (55, 26, 55, 31): 'dtype'}, {}), '(t, dtype)', True, 'import tensorflow.compat.v1 as tf\n'), ((86, 18, 86, 39), 'tensorflow.compat.v1.split', 'tf.split', ({(86, 27, 86, 32): 'state', (86, 34, 86, 35): '2', (86, 37, 86, 38): '1'}, {}), '(state, 2, 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((139, 23, 139, 50), 'tensorflow.compat.v1.square', 'tf.square', ({(139, 33, 139, 49): 'h_reshape - mean'}, {}), '(h_reshape - mean)', True, 'import tensorflow.compat.v1 as tf\n'), ((145, 7, 145, 31), 'tensorflow.compat.v1.variable_scope', 'tf.variable_scope', ({(145, 25, 145, 30): 'scope'}, {}), '(scope)', True, 'import tensorflow.compat.v1 as tf\n'), ((170, 23, 170, 43), 'tensorflow.compat.v1.square', 'tf.square', ({(170, 33, 170, 42): 'x_shifted'}, {}), '(x_shifted)', True, 'import tensorflow.compat.v1 as tf\n'), ((172, 7, 172, 31), 'tensorflow.compat.v1.variable_scope', 'tf.variable_scope', ({(172, 25, 172, 30): 'scope'}, {}), '(scope)', True, 'import tensorflow.compat.v1 as tf\n'), ((207, 7, 207, 43), 'tensorflow.compat.v1.variable_scope', 'tf.variable_scope', ({(207, 25, 207, 42): "(scope or 'linear')"}, {}), "(scope or 'linear')", True, 'import tensorflow.compat.v1 as tf\n'), ((225, 8, 226, 80), 'tensorflow.compat.v1.get_variable', 'tf.get_variable', (), '', True, 'import tensorflow.compat.v1 as tf\n'), ((233, 11, 233, 26), 'tensorflow.compat.v1.matmul', 'tf.matmul', ({(233, 21, 233, 22): 'x', (233, 24, 233, 25): 'w'}, {}), '(x, w)', True, 'import tensorflow.compat.v1 as tf\n'), ((274, 18, 274, 39), 'tensorflow.compat.v1.split', 'tf.split', ({(274, 27, 274, 32): 'state', (274, 34, 274, 35): '2', (274, 37, 274, 38): '1'}, {}), '(state, 2, 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((379, 30, 379, 51), 'tensorflow.compat.v1.split', 'tf.split', ({(379, 39, 379, 44): 'state', (379, 46, 379, 47): '2', (379, 49, 379, 50): '1'}, {}), '(state, 2, 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((91, 13, 91, 34), 'tensorflow.compat.v1.split', 'tf.split', ({(91, 22, 91, 27): 'state', (91, 29, 91, 30): '2', (91, 32, 91, 33): '1'}, {}), '(state, 2, 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((100, 13, 101, 67), 'tensorflow.compat.v1.get_variable', 'tf.get_variable', (), '', True, 'import tensorflow.compat.v1 as tf\n'), ((102, 13, 103, 75), 'tensorflow.compat.v1.get_variable', 'tf.get_variable', (), '', True, 'import tensorflow.compat.v1 as tf\n'), ((108, 15, 108, 35), 'tensorflow.compat.v1.concat', 'tf.concat', ({(108, 25, 108, 31): '[x, h]', (108, 33, 108, 34): '1'}, {}), '([x, h], 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((109, 15, 109, 41), 'tensorflow.compat.v1.concat', 'tf.concat', ({(109, 25, 109, 37): '[w_xh, w_hh]', (109, 39, 109, 40): '0'}, {}), '([w_xh, w_hh], 0)', True, 'import tensorflow.compat.v1 as tf\n'), ((112, 19, 112, 41), 'tensorflow.compat.v1.split', 'tf.split', ({(112, 28, 112, 34): 'hidden', (112, 36, 112, 37): '4', (112, 39, 112, 40): '1'}, {}), '(hidden, 4, 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((217, 15, 217, 43), 'tensorflow.compat.v1.constant_initializer', 'tf.constant_initializer', ({(217, 39, 217, 42): '0.0'}, {}), '(0.0)', True, 'import tensorflow.compat.v1 as tf\n'), ((279, 13, 279, 34), 'tensorflow.compat.v1.split', 'tf.split', ({(279, 22, 279, 27): 'state', (279, 29, 279, 30): '2', (279, 32, 279, 33): '1'}, {}), '(state, 2, 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((289, 13, 290, 67), 'tensorflow.compat.v1.get_variable', 'tf.get_variable', (), '', True, 'import tensorflow.compat.v1 as tf\n'), ((291, 13, 292, 75), 'tensorflow.compat.v1.get_variable', 'tf.get_variable', (), '', True, 'import tensorflow.compat.v1 as tf\n'), ((294, 15, 294, 35), 'tensorflow.compat.v1.concat', 'tf.concat', ({(294, 25, 294, 31): '[x, h]', (294, 33, 294, 34): '1'}, {}), '([x, h], 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((295, 15, 295, 41), 'tensorflow.compat.v1.concat', 'tf.concat', ({(295, 25, 295, 37): '[w_xh, w_hh]', (295, 39, 295, 40): '0'}, {}), '([w_xh, w_hh], 0)', True, 'import tensorflow.compat.v1 as tf\n'), ((296, 15, 296, 40), 'tensorflow.compat.v1.matmul', 'tf.matmul', ({(296, 25, 296, 31): 'concat', (296, 33, 296, 39): 'w_full'}, {}), '(concat, w_full)', True, 'import tensorflow.compat.v1 as tf\n'), ((300, 19, 300, 41), 'tensorflow.compat.v1.split', 'tf.split', ({(300, 28, 300, 34): 'concat', (300, 36, 300, 37): '4', (300, 39, 300, 40): '1'}, {}), '(concat, 4, 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((310, 18, 310, 46), 'tensorflow.compat.v1.concat', 'tf.concat', ({(310, 28, 310, 42): '[new_h, new_c]', (310, 44, 310, 45): '(1)'}, {}), '([new_h, new_c], 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((388, 9, 388, 33), 'tensorflow.compat.v1.variable_scope', 'tf.variable_scope', ({(388, 27, 388, 32): 'scope'}, {}), '(scope)', True, 'import tensorflow.compat.v1 as tf\n'), ((404, 15, 404, 40), 'tensorflow.compat.v1.multiply', 'tf.multiply', ({(404, 27, 404, 32): 'alpha', (404, 34, 404, 39): 'layer'}, {}), '(alpha, layer)', True, 'import tensorflow.compat.v1 as tf\n'), ((426, 25, 426, 46), 'tensorflow.compat.v1.split', 'tf.split', ({(426, 34, 426, 39): 'state', (426, 41, 426, 42): '2', (426, 44, 426, 45): '1'}, {}), '(state, 2, 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((429, 25, 430, 72), 'tensorflow.compat.v1.concat', 'tf.concat', ({(430, 10, 430, 68): '[total_h[:, self.num_units:], total_c[:, self.num_units:]]', (430, 70, 430, 71): '1'}, {}), '([total_h[:, self.num_units:], total_c[:, self.num_units:]], 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((440, 13, 441, 67), 'tensorflow.compat.v1.get_variable', 'tf.get_variable', (), '', True, 'import tensorflow.compat.v1 as tf\n'), ((442, 13, 443, 75), 'tensorflow.compat.v1.get_variable', 'tf.get_variable', (), '', True, 'import tensorflow.compat.v1 as tf\n'), ((449, 20, 449, 40), 'tensorflow.compat.v1.concat', 'tf.concat', ({(449, 30, 449, 36): '[x, h]', (449, 38, 449, 39): '1'}, {}), '([x, h], 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((455, 11, 455, 29), 'tensorflow.compat.v1.matmul', 'tf.matmul', ({(455, 21, 455, 22): 'x', (455, 24, 455, 28): 'w_xh'}, {}), '(x, w_xh)', True, 'import tensorflow.compat.v1 as tf\n'), ((456, 11, 456, 29), 'tensorflow.compat.v1.matmul', 'tf.matmul', ({(456, 21, 456, 22): 'h', (456, 24, 456, 28): 'w_hh'}, {}), '(h, w_hh)', True, 'import tensorflow.compat.v1 as tf\n'), ((459, 23, 459, 41), 'tensorflow.compat.v1.split', 'tf.split', ({(459, 32, 459, 34): 'xh', (459, 36, 459, 37): '4', (459, 39, 459, 40): '1'}, {}), '(xh, 4, 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((466, 23, 466, 41), 'tensorflow.compat.v1.split', 'tf.split', ({(466, 32, 466, 34): 'hh', (466, 36, 466, 37): '4', (466, 39, 466, 40): '1'}, {}), '(hh, 4, 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((473, 23, 473, 43), 'tensorflow.compat.v1.split', 'tf.split', ({(473, 32, 473, 36): 'bias', (473, 38, 473, 39): '4', (473, 41, 473, 42): '0'}, {}), '(bias, 4, 0)', True, 'import tensorflow.compat.v1 as tf\n'), ((494, 25, 494, 56), 'tensorflow.compat.v1.split', 'tf.split', ({(494, 34, 494, 49): 'hyper_new_state', (494, 51, 494, 52): '2', (494, 54, 494, 55): '1'}, {}), '(hyper_new_state, 2, 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((495, 20, 495, 50), 'tensorflow.compat.v1.concat', 'tf.concat', ({(495, 30, 495, 46): '[new_h, hyper_h]', (495, 48, 495, 49): '1'}, {}), '([new_h, hyper_h], 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((496, 20, 496, 50), 'tensorflow.compat.v1.concat', 'tf.concat', ({(496, 30, 496, 46): '[new_c, hyper_c]', (496, 48, 496, 49): '1'}, {}), '([new_c, hyper_c], 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((497, 24, 497, 64), 'tensorflow.compat.v1.concat', 'tf.concat', ({(497, 34, 497, 60): '[new_total_h, new_total_c]', (497, 62, 497, 63): '1'}, {}), '([new_total_h, new_total_c], 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((110, 15, 110, 40), 'tensorflow.compat.v1.matmul', 'tf.matmul', ({(110, 25, 110, 31): 'concat', (110, 33, 110, 39): 'w_full'}, {}), '(concat, w_full)', True, 'import tensorflow.compat.v1 as tf\n'), ((117, 12, 117, 22), 'tensorflow.compat.v1.tanh', 'tf.tanh', ({(117, 20, 117, 21): 'j'}, {}), '(j)', True, 'import tensorflow.compat.v1 as tf\n'), ((120, 14, 120, 28), 'tensorflow.compat.v1.tanh', 'tf.tanh', ({(120, 22, 120, 27): 'new_c'}, {}), '(new_c)', True, 'import tensorflow.compat.v1 as tf\n'), ((120, 31, 120, 44), 'tensorflow.compat.v1.sigmoid', 'tf.sigmoid', ({(120, 42, 120, 43): 'o'}, {}), '(o)', True, 'import tensorflow.compat.v1 as tf\n'), ((122, 20, 122, 48), 'tensorflow.compat.v1.concat', 'tf.concat', ({(122, 30, 122, 44): '[new_c, new_h]', (122, 46, 122, 47): '(1)'}, {}), '([new_c, new_h], 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((150, 20, 150, 56), 'tensorflow.compat.v1.constant_initializer', 'tf.constant_initializer', ({(150, 44, 150, 55): 'gamma_start'}, {}), '(gamma_start)', True, 'import tensorflow.compat.v1 as tf\n'), ((177, 20, 177, 56), 'tensorflow.compat.v1.constant_initializer', 'tf.constant_initializer', ({(177, 44, 177, 55): 'gamma_start'}, {}), '(gamma_start)', True, 'import tensorflow.compat.v1 as tf\n'), ((191, 21, 191, 40), 'tensorflow.compat.v1.square', 'tf.square', ({(191, 31, 191, 39): 'x - mean'}, {}), '(x - mean)', True, 'import tensorflow.compat.v1 as tf\n'), ((219, 15, 219, 52), 'tensorflow.compat.v1.constant_initializer', 'tf.constant_initializer', ({(219, 39, 219, 51): 'weight_start'}, {}), '(weight_start)', True, 'import tensorflow.compat.v1 as tf\n'), ((232, 13, 232, 28), 'tensorflow.compat.v1.matmul', 'tf.matmul', ({(232, 23, 232, 24): 'x', (232, 26, 232, 27): 'w'}, {}), '(x, w)', True, 'import tensorflow.compat.v1 as tf\n'), ((305, 12, 305, 22), 'tensorflow.compat.v1.tanh', 'tf.tanh', ({(305, 20, 305, 21): 'j'}, {}), '(j)', True, 'import tensorflow.compat.v1 as tf\n'), ((308, 59, 308, 72), 'tensorflow.compat.v1.sigmoid', 'tf.sigmoid', ({(308, 70, 308, 71): 'o'}, {}), '(o)', True, 'import tensorflow.compat.v1 as tf\n'), ((482, 17, 482, 43), 'tensorflow.compat.v1.concat', 'tf.concat', ({(482, 27, 482, 39): '[i, j, f, o]', (482, 41, 482, 42): '1'}, {}), '([i, j, f, o], 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((484, 21, 484, 43), 'tensorflow.compat.v1.split', 'tf.split', ({(484, 30, 484, 36): 'concat', (484, 38, 484, 39): '4', (484, 41, 484, 42): '1'}, {}), '(concat, 4, 1)', True, 'import tensorflow.compat.v1 as tf\n'), ((489, 12, 489, 22), 'tensorflow.compat.v1.tanh', 'tf.tanh', ({(489, 20, 489, 21): 'j'}, {}), '(j)', True, 'import tensorflow.compat.v1 as tf\n'), ((492, 67, 492, 80), 'tensorflow.compat.v1.sigmoid', 'tf.sigmoid', ({(492, 78, 492, 79): 'o'}, {}), '(o)', True, 'import tensorflow.compat.v1 as tf\n'), ((106, 22, 106, 50), 'tensorflow.compat.v1.constant_initializer', 'tf.constant_initializer', ({(106, 46, 106, 49): '0.0'}, {}), '(0.0)', True, 'import tensorflow.compat.v1 as tf\n'), ((115, 26, 115, 36), 'tensorflow.compat.v1.tanh', 'tf.tanh', ({(115, 34, 115, 35): 'j'}, {}), '(j)', True, 'import tensorflow.compat.v1 as tf\n'), ((119, 18, 119, 50), 'tensorflow.compat.v1.sigmoid', 'tf.sigmoid', ({(119, 29, 119, 49): '(f + self.forget_bias)'}, {}), '(f + self.forget_bias)', True, 'import tensorflow.compat.v1 as tf\n'), ((119, 53, 119, 66), 'tensorflow.compat.v1.sigmoid', 'tf.sigmoid', ({(119, 64, 119, 65): 'i'}, {}), '(i)', True, 'import tensorflow.compat.v1 as tf\n'), ((147, 6, 147, 29), 'tensorflow.compat.v1.get_variable_scope', 'tf.get_variable_scope', ({}, {}), '()', True, 'import tensorflow.compat.v1 as tf\n'), ((153, 50, 153, 78), 'tensorflow.compat.v1.constant_initializer', 'tf.constant_initializer', ({(153, 74, 153, 77): '0.0'}, {}), '(0.0)', True, 'import tensorflow.compat.v1 as tf\n'), ((174, 6, 174, 29), 'tensorflow.compat.v1.get_variable_scope', 'tf.get_variable_scope', ({}, {}), '()', True, 'import tensorflow.compat.v1 as tf\n'), ((180, 46, 180, 74), 'tensorflow.compat.v1.constant_initializer', 'tf.constant_initializer', ({(180, 70, 180, 73): '0.0'}, {}), '(0.0)', True, 'import tensorflow.compat.v1 as tf\n'), ((209, 6, 209, 29), 'tensorflow.compat.v1.get_variable_scope', 'tf.get_variable_scope', ({}, {}), '()', True, 'import tensorflow.compat.v1 as tf\n'), ((221, 15, 221, 64), 'tensorflow.compat.v1.random_normal_initializer', 'tf.random_normal_initializer', (), '', True, 'import tensorflow.compat.v1 as tf\n'), ((231, 22, 231, 57), 'tensorflow.compat.v1.constant_initializer', 'tf.constant_initializer', ({(231, 46, 231, 56): 'bias_start'}, {}), '(bias_start)', True, 'import tensorflow.compat.v1 as tf\n'), ((303, 26, 303, 36), 'tensorflow.compat.v1.tanh', 'tf.tanh', ({(303, 34, 303, 35): 'j'}, {}), '(j)', True, 'import tensorflow.compat.v1 as tf\n'), ((307, 18, 307, 50), 'tensorflow.compat.v1.sigmoid', 'tf.sigmoid', ({(307, 29, 307, 49): '(f + self.forget_bias)'}, {}), '(f + self.forget_bias)', True, 'import tensorflow.compat.v1 as tf\n'), ((307, 53, 307, 66), 'tensorflow.compat.v1.sigmoid', 'tf.sigmoid', ({(307, 64, 307, 65): 'i'}, {}), '(i)', True, 'import tensorflow.compat.v1 as tf\n'), ((446, 22, 446, 50), 'tensorflow.compat.v1.constant_initializer', 'tf.constant_initializer', ({(446, 46, 446, 49): '0.0'}, {}), '(0.0)', True, 'import tensorflow.compat.v1 as tf\n'), ((487, 26, 487, 36), 'tensorflow.compat.v1.tanh', 'tf.tanh', ({(487, 34, 487, 35): 'j'}, {}), '(j)', True, 'import tensorflow.compat.v1 as tf\n'), ((491, 18, 491, 50), 'tensorflow.compat.v1.sigmoid', 'tf.sigmoid', ({(491, 29, 491, 49): '(f + self.forget_bias)'}, {}), '(f + self.forget_bias)', True, 'import tensorflow.compat.v1 as tf\n'), ((491, 53, 491, 66), 'tensorflow.compat.v1.sigmoid', 'tf.sigmoid', ({(491, 64, 491, 65): 'i'}, {}), '(i)', True, 'import tensorflow.compat.v1 as tf\n')]