code: """
   Copyright (c) 2019 Intel Corporation
   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.
  """
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  from ..box_utils import match, log_sum_exp
  
  
  class MultiBoxLoss(nn.Module):
      """SSD Weighted Loss Function
      Compute Targets:
          1) Produce Confidence Target Indices by matching  ground truth boxes
             with (default) 'priorboxes' that have jaccard index > threshold parameter
             (default threshold: 0.5).
          2) Produce localization target by 'encoding' variance into offsets of ground
             truth boxes and their matched  'priorboxes'.
          3) Hard negative mining to filter the excessive number of negative examples
             that comes with using a large number of default bounding boxes.
             (default negative:positive ratio 3:1)
      Objective Loss:
          L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N
          Where, Lconf is the CrossEntropy Loss and Lloc is the SmoothL1 Loss
          weighted by α which is set to 1 by cross val.
          Args:
              c: class confidences,
              l: predicted boxes,
              g: ground truth boxes
              N: number of matched default boxes
          See: https://arxiv.org/pdf/1512.02325.pdf for more details.
      """
  
      def __init__(self, cfg, num_classes, overlap_thresh, prior_for_matching,
                   bkg_label, neg_mining, neg_pos, neg_overlap, encode_target, device=None):
          super(MultiBoxLoss, self).__init__()
          self.device = device
          self.num_classes = num_classes
          self.threshold = overlap_thresh
          self.background_label = bkg_label
          self.encode_target = encode_target
          self.use_prior_for_matching = prior_for_matching
          self.do_neg_mining = neg_mining
          self.negpos_ratio = neg_pos
          self.neg_overlap = neg_overlap
  
      def forward(self, predictions, targets):
          """Multibox Loss
          Args:
              predictions (tuple): A tuple containing loc preds, conf preds,
              and prior boxes from SSD net.
                  conf shape: torch.size(batch_size,num_priors,num_classes)
                  loc shape: torch.size(batch_size,num_priors,4)
                  priors shape: torch.size(num_priors,4)
  
              ground_truth (tensor): Ground truth boxes and labels for a batch,
                  shape: [batch_size,num_objs,5] (last idx is the label).
          """
          loc_data, conf_data, priors = predictions
          batch = loc_data.size(0)
          num_priors = loc_data.size(1)
  
          # match priors (default boxes) and ground truth boxes
          loc_t = torch.Tensor(batch, num_priors, 4).to(self.device)
          conf_t = torch.LongTensor(batch, num_priors).to(self.device)
          for idx in range(batch):
              truths = targets[idx][:, :-1].data
              labels = targets[idx][:, -1].data
              defaults = priors.data
              match(self.threshold, truths, defaults[0], labels, loc_t, conf_t, idx)
          pos = conf_t > 0
          num_pos = pos.sum(dim=1, keepdim=True)
  
          # Localization Loss (Smooth L1)
          # Shape: [batch,num_priors,4]
          pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data)
          loc_p = loc_data[pos_idx].view(-1, 4)
          loc_t = loc_t[pos_idx].view(-1, 4)
          loss_l = F.smooth_l1_loss(loc_p, loc_t, reduction='sum')
  
          # Compute max conf across batch for hard negative mining
          batch_conf = conf_data.view(-1, self.num_classes)
  
          loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1))
  
          # Hard Negative Mining
          loss_c = loss_c.view(batch, -1)
          loss_c[pos] = 0  # filter out pos boxes for now
          _, loss_idx = loss_c.sort(1, descending=True)
          _, idx_rank = loss_idx.sort(1)
          num_pos = pos.long().sum(1, keepdim=True)
          num_neg = torch.clamp(self.negpos_ratio * num_pos, max=pos.size(1) - 1)
          neg = idx_rank < num_neg.expand_as(idx_rank)
  
          # Confidence Loss Including Positive and Negative Examples
          pos_idx = pos.unsqueeze(2).expand_as(conf_data)
          neg_idx = neg.unsqueeze(2).expand_as(conf_data)
          conf_p = conf_data[(pos_idx + neg_idx).gt(0)].view(-1, self.num_classes)
          targets_weighted = conf_t[(pos + neg).gt(0)]
          loss_c = F.cross_entropy(conf_p, targets_weighted, reduction='sum')
  
          # Sum of losses: L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N
  
          N = num_pos.data.sum().to(torch.float)
          loss_l /= N
          loss_c /= N
          return loss_l, loss_c
   # Copyright (C) 2019 Intel Corporation
  #
  # 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.
  
  """ This module contains architecture of Text Recognition model."""
  
  import tensorflow as tf
  from tensorflow.contrib import rnn
  import tensorflow.contrib.slim as slim
  
  
  class TextRecognition:
      """ Text recognition model definition. """
  
      def __init__(self, is_training, num_classes, backbone_dropout=0.0):
          self.is_training = is_training
          self.lstm_dim = 256
          self.num_classes = num_classes
          self.backbone_dropout = backbone_dropout
  
      def __call__(self, inputdata):
          with tf.variable_scope('shadow'):
              features = self.feature_extractor(inputdata=inputdata)
              logits = self.encoder_decoder(inputdata=tf.squeeze(features, axis=1))
  
          return logits
  
      # pylint: disable=too-many-locals
      def feature_extractor(self, inputdata):
          """ Extracts features from input text image. """
  
          with slim.arg_scope([slim.conv2d], padding='SAME',
                              weights_initializer=tf.contrib.layers.variance_scaling_initializer(),
                              weights_regularizer=slim.l2_regularizer(0.00025),
                              biases_initializer=None, activation_fn=None):
              with slim.arg_scope([slim.batch_norm], updates_collections=None):
                  bn0 = slim.batch_norm(inputdata, 0.9, scale=True, is_training=self.is_training,
                                        activation_fn=None)
  
                  dropout1 = slim.dropout(bn0, keep_prob=1.0 - self.backbone_dropout,
                                          is_training=self.is_training)
                  conv1 = slim.conv2d(dropout1, num_outputs=64, kernel_size=3)
                  bn1 = slim.batch_norm(conv1, 0.9, scale=True, is_training=self.is_training,
                                        activation_fn=tf.nn.relu)
                  pool1 = slim.max_pool2d(bn1, kernel_size=2, stride=2)
  
                  dropout2 = slim.dropout(pool1, keep_prob=1.0 - self.backbone_dropout,
                                          is_training=self.is_training)
                  conv2 = slim.conv2d(dropout2, num_outputs=128, kernel_size=3)
                  bn2 = slim.batch_norm(conv2, 0.9, scale=True, is_training=self.is_training,
                                        activation_fn=tf.nn.relu)
                  pool2 = slim.max_pool2d(bn2, kernel_size=2, stride=2)
  
                  dropout3 = slim.dropout(pool2, keep_prob=1.0 - self.backbone_dropout,
                                          is_training=self.is_training)
                  conv3 = slim.conv2d(dropout3, num_outputs=256, kernel_size=3)
                  bn3 = slim.batch_norm(conv3, 0.9, scale=True, is_training=self.is_training,
                                        activation_fn=tf.nn.relu)
  
                  dropout4 = slim.dropout(bn3, keep_prob=1.0 - self.backbone_dropout,
                                          is_training=self.is_training)
                  conv4 = slim.conv2d(dropout4, num_outputs=256, kernel_size=3)
                  bn4 = slim.batch_norm(conv4, 0.9, scale=True, is_training=self.is_training,
                                        activation_fn=tf.nn.relu)
                  pool4 = slim.max_pool2d(bn4, kernel_size=[2, 1], stride=[2, 1])
  
                  dropout5 = slim.dropout(pool4, keep_prob=1.0 - self.backbone_dropout,
                                          is_training=self.is_training)
                  conv5 = slim.conv2d(dropout5, num_outputs=512, kernel_size=3)
                  bn5 = slim.batch_norm(conv5, 0.9, scale=True, is_training=self.is_training,
                                        activation_fn=tf.nn.relu)
  
                  dropout6 = slim.dropout(bn5, keep_prob=1.0 - self.backbone_dropout,
                                          is_training=self.is_training)
                  conv6 = slim.conv2d(dropout6, num_outputs=512, kernel_size=3)
                  bn6 = slim.batch_norm(conv6, 0.9, scale=True, is_training=self.is_training,
                                        activation_fn=tf.nn.relu)
                  pool6 = slim.max_pool2d(bn6, kernel_size=[2, 1], stride=[2, 1])
  
                  dropout7 = slim.dropout(pool6, keep_prob=1.0 - self.backbone_dropout,
                                          is_training=self.is_training)
                  conv7 = slim.conv2d(dropout7, num_outputs=512, kernel_size=2, stride=[2, 1])
                  bn7 = slim.batch_norm(conv7, 0.9, scale=True, is_training=self.is_training,
                                        activation_fn=tf.nn.relu)
  
          return bn7
  
      def encoder_decoder(self, inputdata):
          """ LSTM-based encoder-decoder module. """
  
          with tf.variable_scope('LSTMLayers'):
              [batch_size, width, _] = inputdata.get_shape().as_list()
  
              with tf.variable_scope('encoder'):
                  forward_cells = []
                  backward_cells = []
  
                  for _ in range(2):
                      forward_cells.append(tf.nn.rnn_cell.LSTMCell(self.lstm_dim))
                      backward_cells.append(tf.nn.rnn_cell.LSTMCell(self.lstm_dim))
  
                  encoder_layer, _, _ = rnn.stack_bidirectional_dynamic_rnn(
                      forward_cells, backward_cells, inputdata, dtype=tf.float32)
  
              with tf.variable_scope('decoder'):
                  forward_cells = []
                  backward_cells = []
  
                  for _ in range(2):
                      forward_cells.append(tf.nn.rnn_cell.LSTMCell(self.lstm_dim))
                      backward_cells.append(tf.nn.rnn_cell.LSTMCell(self.lstm_dim))
  
                  decoder_layer, _, _ = rnn.stack_bidirectional_dynamic_rnn(
                      forward_cells, backward_cells, encoder_layer, dtype=tf.float32)
  
              rnn_reshaped = tf.reshape(decoder_layer, [batch_size * width, -1])
  
              logits = slim.fully_connected(rnn_reshaped, self.num_classes, activation_fn=None)
              logits = tf.reshape(logits, [batch_size, width, self.num_classes])
              rnn_out = tf.transpose(logits, (1, 0, 2))
  
          return rnn_out
   """
   Copyright (c) 2019-2020 Intel Corporation
   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.
  """
  import logging
  import sys
  from collections import namedtuple
  from typing import List, Tuple
  
  import copy
  import onnx
  import os
  import pytest
  import torch
  import torch.nn.functional as F
  from functools import partial
  from torch import nn
  from torch.autograd import Variable
  from torch.nn.utils.rnn import PackedSequence
  
  from nncf.dynamic_graph.context import TracingContext
  from nncf.dynamic_graph.transform_graph import replace_modules
  from nncf.model_creation import create_compressed_model
  from nncf.layers import LSTMCellNNCF, NNCF_RNN, ITERATION_MODULES
  from tests.modules.seq2seq.gnmt import GNMT
  from tests.test_helpers import get_empty_config, get_grads, create_compressed_model_and_algo_for_test
  
  
  def replace_lstm(model):
      def replace_fn(module_):
          if not isinstance(module_, nn.LSTM):
              return module_
          device = next(module_.parameters()).device
          custom_lstm = NNCF_RNN('LSTM', input_size=module_.input_size, hidden_size=module_.hidden_size,
                                 num_layers=module_.num_layers, bidirectional=module_.bidirectional,
                                 batch_first=module_.batch_first, dropout=module_.dropout,
                                 bias=module_.bias)
  
          def get_param_names(bias):
              # type: (bool) -> List[str]
              suffixes = ['ih', 'hh']
              names = ['weight_' + suffix for suffix in suffixes]
              if bias:
                  names += ['bias_' + suffix for suffix in suffixes]
              return names
  
          for l in range(custom_lstm.num_layers):
              for d in range(custom_lstm.num_directions):
                  for name in get_param_names(custom_lstm.bias):
                      suffix = '_reverse' if d == 1 else ''
                      param_name = name + '_l{}{}'.format(l, suffix)
                      param = getattr(module_, param_name)
                      getattr(custom_lstm, param_name).data.copy_(param.data)
          custom_lstm.to(device)
          return custom_lstm
  
      if isinstance(model, nn.LSTM):
          return replace_fn(model)
      affected_scopes = []
      return replace_modules(model, replace_fn, affected_scopes)[0]
  
  def clone_test_data(data_list):
      # type: (LSTMTestData) -> List[torch.Tensor]
      results = []
      x = data_list[0]
      result = x if isinstance(x, PackedSequence) else x.clone()
      results.append(result)
      for tensor_list in data_list[1:]:
          result = ()
          for tensor in tensor_list:
              if isinstance(tensor, Variable):
                  sub_result = tensor.data.clone()
                  sub_result = Variable(sub_result, requires_grad=True)
              else:
                  sub_result = tensor.clone()
              result += (sub_result,)
          results.append(result)
      return results
  
  
  LSTMTestSizes = namedtuple('LSTMTestSizes', ['input_size', 'hidden_size', 'batch', 'seq_length'])
  LSTMTestData = namedtuple('LSTMTestData', ['x', 'h0', 'c0', 'weight_ih', 'weight_hh', 'bias_ih', 'bias_hh'])
  
  
  @pytest.mark.parametrize('sizes',
                           [LSTMTestSizes(512, 768, 128, 50),
                            LSTMTestSizes(3, 3, 3, 3),
                            LSTMTestSizes(1, 1, 1, 1)], ids=lambda val: '[{}]'.format('-'.join([str(v) for v in val])))
  class TestLSTMCell:
      @staticmethod
      def generate_lstm_data(p, num_layers=1, num_directions=1, variable_length=False, sorted_=True, batch_first=True,
                             is_cuda=False, bias=True, empty_initial=False, is_backward=False):
          # type: (LSTMTestSizes, int, int, bool, bool, bool, bool, bool, bool, bool) -> LSTMTestData
          num_chunks = 4
          seq_list = []
          if variable_length:
              seq_lens = torch.IntTensor(p.batch).random_(1, p.seq_length + 1)
              if sorted_:
                  seq_lens = torch.sort(seq_lens, descending=True).values
              for seq_size in seq_lens:
                  seq_list.append(torch.randn(seq_size.item(), p.input_size))
              padded_seq_batch = torch.nn.utils.rnn.pad_sequence(seq_list, batch_first=batch_first)
              x_data = torch.nn.utils.rnn.pack_padded_sequence(padded_seq_batch, lengths=seq_lens,
                                                               batch_first=batch_first, enforce_sorted=sorted_)
  
          else:
              size = (p.seq_length, p.batch, p.input_size)
              if batch_first:
                  size = (p.batch, p.seq_length, p.input_size)
              x_data = torch.randn(*size)
  
          def wrap_tensor(tensor):
              wrapped = tensor
              if is_cuda:
                  wrapped = wrapped.cuda()
              if is_backward:
                  wrapped = Variable(wrapped, requires_grad=True)
              return wrapped
  
          if is_cuda:
              x_data = x_data.cuda()
          h0, c0, wih, whh, bih, bhh = ([] for _ in range(6))
          for layer_ in range(num_layers):
              for _ in range(num_directions):
                  layer_input_size = p.input_size if layer_ == 0 else p.hidden_size * num_directions
                  if not empty_initial:
                      h0.append(wrap_tensor(torch.randn(p.batch, p.hidden_size)))
                      c0.append(wrap_tensor(torch.randn(p.batch, p.hidden_size)))
                  wih.append(wrap_tensor(torch.rand(num_chunks * p.hidden_size, layer_input_size)))
                  whh.append(wrap_tensor(torch.rand(num_chunks * p.hidden_size, p.hidden_size)))
                  if bias:
                      bih.append(wrap_tensor(torch.rand(num_chunks * p.hidden_size)))
                      bhh.append(wrap_tensor(torch.rand(num_chunks * p.hidden_size)))
          result = LSTMTestData(x_data, h0, c0, wih, whh, bih, bhh)
          return result
  
      @staticmethod
      def set_weights(cell, data):
          # type: (nn.LSTMCell, LSTMTestData) -> None
          for name in TestLSTM.get_param_names(bias=True):
              param = getattr(data, name)
              if param:
                  getattr(cell, name).data.copy_(param[0].data)
  
      def test_forward_lstm_cell(self, sizes, _seed):
          p = sizes
          ref_data = TestLSTMCell.generate_lstm_data(p, batch_first=False)
          test_data = LSTMTestData(*clone_test_data(ref_data))
  
          ref_rnn = nn.LSTMCell(p.input_size, p.hidden_size)
          TestLSTMCell.set_weights(ref_rnn, ref_data)
          test_rnn = LSTMCellNNCF(p.input_size, p.hidden_size)
          TestLSTMCell.set_weights(test_rnn, test_data)
  
          for i in range(p.seq_length):
              ref_result = ref_rnn(ref_data.x[i], (ref_data.h0[0], ref_data.c0[0]))
              test_result = test_rnn(test_data.x[i], (test_data.h0[0], test_data.c0[0]))
              for (ref, test) in list(zip(ref_result, test_result)):
                  torch.testing.assert_allclose(test, ref)
  
      def test_backward_lstm_cell(self, sizes, _seed):
          p = sizes
          ref_data = TestLSTMCell.generate_lstm_data(p, batch_first=False, is_backward=True)
          with torch.no_grad():
              test_data = LSTMTestData(*clone_test_data(ref_data))
  
          ref_rnn = nn.LSTMCell(p.input_size, p.hidden_size)
          TestLSTMCell.set_weights(ref_rnn, ref_data)
          test_rnn = LSTMCellNNCF(p.input_size, p.hidden_size)
          TestLSTMCell.set_weights(test_rnn, test_data)
  
          for i in range(p.seq_length):
              ref_result = ref_rnn(ref_data.x[i], (ref_data.h0[0], ref_data.c0[0]))
              test_result = test_rnn(test_data.x[i], (test_data.h0[0], test_data.c0[0]))
              ref_result[0].sum().backward()
              test_result[0].sum().backward()
              ref_grads = get_grads([ref_data.h0[0], ref_data.c0[0]])
              ref_grads += get_grads([ref_rnn.weight_ih, ref_rnn.weight_hh, ref_rnn.bias_ih, ref_rnn.bias_hh])
              test_grads = get_grads([ref_data.h0[0], ref_data.c0[0]])
              test_grads += get_grads([test_rnn.weight_ih, test_rnn.weight_hh, test_rnn.bias_ih, test_rnn.bias_hh])
              for (ref, test) in list(zip(test_grads, ref_grads)):
                  torch.testing.assert_allclose(test, ref)
  
  
  def test_export_lstm_cell(tmp_path):
      config = get_empty_config(model_size=1, input_sample_size=(1, 1))
      config['compression'] = {'algorithm': 'quantization'}
  
      model, algo = create_compressed_model_and_algo_for_test(LSTMCellNNCF(1, 1), config)
  
      test_path = str(tmp_path.joinpath('test.onnx'))
      algo.export_model(test_path)
      assert os.path.exists(test_path)
  
      onnx_num = 0
      model = onnx.load(test_path)
      # pylint: disable=no-member
      for node in model.graph.node:
          if node.op_type == 'FakeQuantize':
              onnx_num += 1
      assert onnx_num == 12
  
  
  @pytest.mark.parametrize('sizes',
                           [LSTMTestSizes(512, 324, 128, 50),
                            LSTMTestSizes(3, 3, 3, 3),
                            LSTMTestSizes(1, 1, 1, 1)], ids=lambda val: '[{}]'.format('-'.join([str(v) for v in val])))
  @pytest.mark.parametrize('bidirectional', (True, False), ids=('bi', 'uni'))
  @pytest.mark.parametrize("bias", [True, False], ids=['bias', 'no_bias'])
  @pytest.mark.parametrize('num_layers', [1, 2], ids=['single_layer', 'stacked'])
  @pytest.mark.parametrize('batch_first', [True, False], ids=['batch_first', 'seq_first'])
  @pytest.mark.parametrize(('variable_length', 'sorted_'),
                           ([True, True],
                            [True, False],
                            [False, False]), ids=['packed_sorted', 'packed_unsorted', 'not_packed'])
  @pytest.mark.parametrize('is_cuda', [True, False], ids=['cuda', 'cpu'])
  @pytest.mark.parametrize('empty_initial', [True, False], ids=['no_initial', 'with_initial'])
  # TODO: dropout gives different result. Looks like different random seed on CPU
  # @pytest.mark.parametrize('dropout', [0, 0.9], ids=['no_dropout', 'with_dropout'])
  @pytest.mark.parametrize('dropout', [0], ids=['no_dropout'])
  class TestLSTM:
      def test_forward_lstm(self, sizes, bidirectional, num_layers, bias, batch_first, variable_length, sorted_, is_cuda,
                            empty_initial, dropout, _seed):
          num_directions = 2 if bidirectional else 1
          p = sizes
  
          ref_data = TestLSTMCell.generate_lstm_data(p, num_layers, num_directions, variable_length, sorted_, batch_first,
                                                     is_cuda, bias, empty_initial)
  
          ref_rnn = nn.LSTM(input_size=p.input_size, hidden_size=p.hidden_size, num_layers=num_layers,
                            bidirectional=bidirectional, batch_first=batch_first, bias=bias, dropout=dropout)
          self.set_ref_lstm_weights(ref_data, ref_rnn, num_layers, num_directions, bias)
          ref_hidden = None if empty_initial else self.get_ref_lstm_hidden(ref_data)
  
          test_data = LSTMTestData(*clone_test_data(ref_data))
  
          class ModelWrapper(nn.Module):
              def __init__(self, lstm):
                  super().__init__()
                  self.lstm = lstm
  
              def forward(self, *input_):
                  return self.lstm(*input_)
  
          wrapped_ref_rnn = ModelWrapper(ref_rnn)
          wrapped_test_rnn = replace_lstm(copy.deepcopy(wrapped_ref_rnn))
          test_rnn = wrapped_test_rnn.lstm
          test_hidden = None if empty_initial else self.get_test_lstm_hidden(test_data)
  
          if is_cuda:
              ref_rnn.cuda()
              test_rnn.cuda()
          ref_output, (ref_hn, ref_cn) = ref_rnn(ref_data.x, ref_hidden)
          test_output, (test_hn, test_cn) = test_rnn(test_data.x, test_hidden)
  
          torch.testing.assert_allclose(test_hn[0], ref_hn[0], rtol=1e-3, atol=1e-4)
          torch.testing.assert_allclose(test_cn[0], ref_cn[0], rtol=1e-3, atol=1e-4)
          if variable_length:
              torch.testing.assert_allclose(test_output.batch_sizes, ref_output.batch_sizes)
              torch.testing.assert_allclose(test_output.data, ref_output.data, rtol=1e-2, atol=1e-3)
              if not sorted_:
                  torch.testing.assert_allclose(test_output.sorted_indices, ref_output.sorted_indices)
                  torch.testing.assert_allclose(test_output.unsorted_indices, ref_output.unsorted_indices)
          else:
              torch.testing.assert_allclose(test_output, ref_output, rtol=1e-2, atol=1e-3)
  
      def test_backward_lstm(self, sizes, bidirectional, num_layers, bias, batch_first, variable_length, sorted_, is_cuda,
                             empty_initial, dropout, _seed):
  
          num_directions = 2 if bidirectional else 1
  
          p = sizes
  
          ref_data = TestLSTMCell.generate_lstm_data(p, num_layers, num_directions, variable_length, sorted_, batch_first,
                                                     is_cuda, bias, empty_initial, True)
  
          ref_rnn = nn.LSTM(input_size=p.input_size, hidden_size=p.hidden_size, num_layers=num_layers,
                            bidirectional=bidirectional, batch_first=batch_first, bias=bias, dropout=dropout)
          self.set_ref_lstm_weights(ref_data, ref_rnn, num_layers, num_directions, bias)
          ref_hidden = None if empty_initial else self.get_ref_lstm_hidden(ref_data)
  
          test_data = LSTMTestData(*clone_test_data(ref_data))
          test_rnn = replace_lstm(copy.deepcopy(ref_rnn))
          test_hidden = None if empty_initial else self.get_test_lstm_hidden(test_data)
  
          if is_cuda:
              ref_rnn.cuda()
              test_rnn.cuda()
  
          ref_output, _ = ref_rnn(ref_data.x, ref_hidden)
          test_output, _ = test_rnn(test_data.x, test_hidden)
  
          ref_output[0].sum().backward()
          test_output[0].sum().backward()
  
          ref_grads = get_grads(self.flatten_nested_lists(ref_rnn.all_weights))
          test_grads = get_grads(self.flatten_nested_lists(test_rnn.all_weights))
          if not empty_initial:
              # TODO: compare gradient of all hidden
              ref_grads += get_grads([ref_data.h0[0], ref_data.c0[0]])
              test_grads += get_grads([test_hidden[0][0], test_hidden[1][0]])
          for (ref, test) in list(zip(test_grads, ref_grads)):
              torch.testing.assert_allclose(test, ref, rtol=1e-1, atol=1e-1)
  
      @classmethod
      def flatten_nested_lists(cls, nested_list):
          # type: (List) -> List[torch.Tensor]
          return [tensor for tensor_tuple in nested_list for tensor in tensor_tuple]
  
      @classmethod
      def get_test_lstm_hidden(cls, data):
          # type: (LSTMTestData) -> List[Tuple[torch.Tensor, ...]]
          result = []
          hidden_names = ['h0', 'c0']
          for name in hidden_names:
              hidden_list = getattr(data, name)
              element = ()
              num_hidden = len(hidden_list)
              for i in range(num_hidden):
                  element += (hidden_list[i],)
              result.append(element)
          return result
  
      @classmethod
      def get_ref_lstm_hidden(cls, data):
          # type: (LSTMTestData) -> Tuple[torch.Tensor, torch.Tensor]
          hidden = cls.get_test_lstm_hidden(data)
          hidden_states = [torch.unsqueeze(tensor, dim=0) for tensor in hidden[0]]
          cell_states = [torch.unsqueeze(tensor, dim=0) for tensor in hidden[1]]
          return (
              torch.cat(hidden_states, dim=0),
              torch.cat(cell_states, dim=0)
          )
  
      @classmethod
      def set_ref_lstm_weights(cls, data, nn_lstm, num_layers, num_directions, bias):
          # type: (LSTMTestData, nn.LSTM, int, int, bool) -> None
          for l in range(num_layers):
              for d in range(num_directions):
                  i = l * num_directions + d
                  for name in cls.get_param_names(bias):
                      suffix = '_reverse' if d == 1 else ''
                      param = getattr(data, name)
                      param_name = name + '_l{}{}'.format(l, suffix)
                      getattr(nn_lstm, param_name).data.copy_(param[i].data)
  
      @classmethod
      def get_param_names(cls, bias):
          # type: (bool) -> List[str]
          suffixes = ['ih', 'hh']
          names = ['weight_' + suffix for suffix in suffixes]
          if bias:
              names += ['bias_' + suffix for suffix in suffixes]
          return names
  
  
  def test_export_stacked_bi_lstm(tmp_path):
      p = LSTMTestSizes(3, 3, 3, 3)
      config = get_empty_config(input_sample_size=(1, p.hidden_size, p.input_size))
      config['compression'] = {'algorithm': 'quantization'}
  
      # TODO: batch_first=True fails with building graph: ambiguous call to mul or sigmoid
      test_rnn = NNCF_RNN('LSTM', input_size=p.input_size, hidden_size=p.hidden_size, num_layers=2, bidirectional=True,
                          batch_first=False)
      model, algo = create_compressed_model_and_algo_for_test(test_rnn, config)
  
      test_path = str(tmp_path.joinpath('test.onnx'))
      algo.export_model(test_path)
      assert os.path.exists(test_path)
  
      onnx_num = 0
      model = onnx.load(test_path)
      # pylint: disable=no-member
      for node in model.graph.node:
          if node.op_type == 'FakeQuantize':
              onnx_num += 1
      assert onnx_num == 50
  
  
  class TestNumberOfNodes:
      logging.basicConfig(level=logging.INFO, stream=sys.stdout)
  
      def test_number_of_calling_fq_for_lstm(self):
          p = LSTMTestSizes(1, 1, 1, 5)
          num_layers = 2
          bidirectional = True
          num_directions = 2 if bidirectional else 1
          bias = True
          batch_first = False
          config = get_empty_config(input_sample_size=(p.seq_length, p.batch, p.input_size))
          config['compression'] = {'algorithm': 'quantization', 'quantize_inputs': True}
  
          test_data = TestLSTMCell.generate_lstm_data(p, num_layers, num_directions, bias=bias, batch_first=batch_first)
  
          test_rnn = NNCF_RNN('LSTM', input_size=p.input_size, hidden_size=p.hidden_size, num_layers=num_layers,
                              bidirectional=bidirectional, bias=bias, batch_first=batch_first)
          TestLSTM.set_ref_lstm_weights(test_data, test_rnn, num_layers, num_directions, bias)
          test_hidden = TestLSTM.get_test_lstm_hidden(test_data)
  
          model, algo = create_compressed_model_and_algo_for_test(test_rnn, config)
  
          class Counter:
              def __init__(self):
                  self.count = 0
  
              def next(self):
                  self.count += 1
  
          def hook(model, input_, counter):
              counter.next()
  
          counters = {}
          for name, quantizer in algo.all_quantizations.items():
              counter = Counter()
              counters[name] = counter
              quantizer.register_forward_pre_hook(partial(hook, counter=counter))
          _ = model(test_data.x, test_hidden)
          assert model.get_graph().get_nodes_count() == 107  # NB: may always fail in debug due to superfluous 'cat' nodes
          assert len(counters) == 50
          for counter in counters.values():
              assert counter.count == p.seq_length
  
      def test_number_of_calling_fq_for_gnmt(self):
          torch.cuda.set_device(0)
          device = torch.device('cuda')
          batch_first = False
          vocab_size = 32000
          model_config = {'hidden_size': 100,
                          'vocab_size': vocab_size,
                          'num_layers': 4,
                          'dropout': 0.2,
                          'batch_first': batch_first,
                          'share_embedding': True,
                          }
          batch_size = 128
          sequence_size = 50
          input_sample_size = (batch_size, sequence_size) if batch_first else (sequence_size, batch_size)
          config = get_empty_config(input_sample_size=input_sample_size)
          config['compression'] = \
              {'algorithm': 'quantization',
               'quantize_inputs': True,
               'quantizable_subgraph_patterns': [["linear", "__add__"],
                                                 ["sigmoid", "__mul__", "__add__"],
                                                 ["__add__", "tanh", "__mul__"],
                                                 ["sigmoid", "__mul__"]],
               'disable_function_quantization_hooks': True}
          config['scopes_without_shape_matching'] = \
              ['GNMT/ResidualRecurrentDecoder[decoder]/RecurrentAttention[att_rnn]/BahdanauAttention[attn]', ]
  
          model = GNMT(**model_config)
          model = replace_lstm(model)
          model.to(device)
  
          def dummy_forward_fn(model, seq_len=sequence_size):
              def gen_packed_sequence():
                  seq_list = []
                  seq_lens = torch.LongTensor(batch_size).random_(1, seq_len + 1)
                  seq_lens = torch.sort(seq_lens, descending=True).values
                  for seq_size in seq_lens:
                      seq_list.append(torch.LongTensor(seq_size.item()).random_(1, vocab_size).to(device))
                  padded_seq_batch = torch.nn.utils.rnn.pad_sequence(seq_list, batch_first=batch_first)
                  return padded_seq_batch, seq_lens
  
              x_data, seq_lens = gen_packed_sequence()
              input_encoder = x_data
              input_enc_len = seq_lens.to(device)
              input_decoder = gen_packed_sequence()[0]
              model(input_encoder, input_enc_len, input_decoder)
  
          algo, model = create_compressed_model(model, config, dummy_forward_fn, dump_graphs=False)
          model.to(device)
  
          class Counter:
              def __init__(self):
                  self.count = 0
  
              def next(self):
                  self.count += 1
  
          def hook(model, input_, counter):
              counter.next()
  
          counters = {}
          for name, quantizer in algo.all_quantizations.items():
              counter = Counter()
              counters[str(name)] = counter
              quantizer.register_forward_pre_hook(partial(hook, counter=counter))
          dummy_forward_fn(model)
          assert model.get_graph().get_nodes_count() == 230  # NB: may always fail in debug due to superfluous 'cat' nodes
          assert len(counters) == 55
          for name, counter in counters.items():
              if 'cell' in name or "LSTMCellForwardNNCF" in name:
                  assert counter.count == sequence_size, name
              else:
                  assert counter.count == 1, name
          new_seq_len = int(sequence_size / 2)
          dummy_forward_fn(model, new_seq_len)
          assert model.get_graph().get_nodes_count() == 230  # NB: may always fail in debug due to superfluous 'cat' nodes
          assert len(counters) == 55
          for name, counter in counters.items():
              if 'cell' in name or "LSTMCellForwardNNCF" in name:
                  assert counter.count == sequence_size + new_seq_len, name
              else:
                  assert counter.count == 2, name
  
      def test_number_of_nodes_for_module_in_loop(self):
          num_iter = 5
  
          class LoopModule(nn.Module):
              @ITERATION_MODULES.register('Inner')
              class Inner(nn.Module):
                  def __init__(self):
                      super().__init__()
                      self.operator1 = torch.sigmoid
                      self.operator2 = torch.tanh
  
                  def forward(self, x):
                      s = self.operator1(x)
                      t = self.operator2(x)
                      result = t + s
                      return result
  
                  @staticmethod
                  def nodes_number():
                      return 3
  
              def __init__(self):
                  super().__init__()
                  self.inner = self.Inner()
  
              def forward(self, x):
                  for _ in range(num_iter):
                      x = self.inner(x)
                  return x
  
              def nodes_number(self):
                  return self.inner.nodes_number()
  
          test_module = LoopModule()
          context = TracingContext()
          with context as ctx:
              _ = test_module(torch.zeros(1))
              assert ctx.graph.get_nodes_count() == test_module.nodes_number()
  
      def test_number_of_nodes_for_module_in_loop__not_input_node(self):
          num_iter = 5
  
          class LoopModule(nn.Module):
              class Inner(nn.Module):
                  def forward(self, x):
                      s = F.sigmoid(x)
                      t = F.tanh(x)
                      result = F.sigmoid(x) * t + F.tanh(x) * s
                      return result
  
                  @staticmethod
                  def nodes_number():
                      return 7
  
              def __init__(self):
                  super().__init__()
                  self.inner = self.Inner()
  
              def forward(self, x):
                  for _ in range(num_iter):
                      x = self.inner(F.relu(x))
                  return x
  
              def nodes_number(self):
                  return self.inner.nodes_number() + num_iter
  
          test_module = LoopModule()
          context = TracingContext()
          with context as ctx:
              _ = test_module(torch.zeros(1))
              assert ctx.graph.get_nodes_count() == test_module.nodes_number()
  
      def test_number_of_nodes_for_module_with_nested_loops(self):
          num_iter = 5
  
          class TestIterModule(nn.Module):
              @ITERATION_MODULES.register()
              class TestIterModule_ResetPoint(nn.Module):
                  def __init__(self, loop_module):
                      super().__init__()
                      self.loop_module = loop_module
  
                  def forward(self, x):
                      return self.loop_module(F.relu(x))
  
              def __init__(self):
                  super().__init__()
                  self.loop_module = self.LoopModule2()
                  self.reset_point = self.TestIterModule_ResetPoint(self.loop_module)
  
              def forward(self, x):
                  for _ in range(num_iter):
                      x = self.reset_point(x)
                  return x
  
              class LoopModule2(nn.Module):
  
                  @ITERATION_MODULES.register()
                  class LoopModule2_ResetPoint(nn.Module):
                      def __init__(self, inner):
                          super().__init__()
                          self.inner = inner
  
                      def forward(self, x):
                          return self.inner(F.relu(x))
  
                  def __init__(self):
                      super().__init__()
                      self.inner = self.Inner()
                      self.reset_helper = self.LoopModule2_ResetPoint(self.inner)
  
                  def forward(self, x):
                      for _ in range(num_iter):
                          self.reset_helper(x)
                      return x
  
                  class Inner(nn.Module):
                      def forward(self, x):
                          s = F.sigmoid(x)
                          t = F.tanh(x)
                          result = t + s
                          return result
  
          test_module = TestIterModule()
          context = TracingContext()
          with context as ctx:
              _ = test_module(torch.zeros(1))
              assert ctx.graph.get_nodes_count() == num_iter
  
      def test_number_of_nodes_for_repeated_module(self):
  
          class LoopModule(nn.Module):
              def __init__(self):
                  super().__init__()
                  self.operator = F.relu
                  self.layers = nn.ModuleList([
                      nn.Conv2d(1, 1, 1),
                      nn.Conv2d(1, 1, 1)
                  ])
  
              def forward(self, x):
                  for layer in self.layers:
                      x = F.relu(layer(x))
                  return x
  
          test_module = LoopModule()
          context = TracingContext()
          with context as ctx:
              x = test_module(torch.zeros(1, 1, 1, 1))
              assert ctx.graph.get_nodes_count() == 4  # NB: may always fail in debug due to superfluous 'cat' nodes
              _ = test_module(x)
              assert ctx.graph.get_nodes_count() == 8  # NB: may always fail in debug due to superfluous 'cat' nodes
  
apis: torch.nn.functional.cross_entropy torch.LongTensor torch.Tensor torch.nn.functional.smooth_l1_loss
apis: tensorflow.contrib.rnn.stack_bidirectional_dynamic_rnn tensorflow.reshape tensorflow.contrib.slim.batch_norm tensorflow.contrib.slim.fully_connected tensorflow.contrib.slim.conv2d tensorflow.contrib.slim.l2_regularizer tensorflow.contrib.slim.dropout tensorflow.variable_scope tensorflow.nn.rnn_cell.LSTMCell tensorflow.contrib.slim.arg_scope tensorflow.squeeze tensorflow.transpose tensorflow.contrib.slim.max_pool2d tensorflow.contrib.layers.variance_scaling_initializer
apis: torch.rand torch.no_grad torch.nn.Conv2d torch.cat torch.nn.functional.sigmoid torch.testing.assert_allclose torch.randn torch.autograd.Variable torch.device torch.sort torch.cuda.set_device torch.unsqueeze torch.nn.utils.rnn.pad_sequence torch.nn.LSTM torch.nn.functional.tanh torch.IntTensor torch.nn.utils.rnn.pack_padded_sequence torch.nn.functional.relu torch.nn.LSTMCell torch.zeros torch.LongTensor