def main():
    tf.set_random_seed(10)
    with tf.Session() as sess:
        rnn_cell = tf.nn.rnn_cell.LSTMCell(10)

        # defining initial state
        initial_state = rnn_cell.zero_state(4, dtype=tf.float32)

        inputs = tf.Variable(tf.random_uniform(shape = (4, 30, 100)), name='input')
        inputs = tf.identity(inputs, "input_node")

        # 'state' is a tensor of shape [batch_size, cell_state_size]
        outputs, state = tf.nn.dynamic_rnn(rnn_cell, inputs, initial_state=initial_state, dtype=tf.float32)

        y1 = tf.identity(outputs, 'outputs')
        y2 = tf.identity(state, 'state')

        t1 = tf.ones([4, 30, 10])
        t2 = tf.ones([4, 10])

        loss = tf.reduce_sum((y1 - t1) * (y1 - t1)) + tf.reduce_sum((y2 - t2) * (y2 - t2))
        tf.identity(loss, name = "lstm_loss")
        # tf.summary.FileWriter('/tmp/log', tf.get_default_graph())

        net_outputs = map(lambda x: tf.get_default_graph().get_tensor_by_name(x), argv[2].split(','))
        run_model(net_outputs, argv[1], None, argv[3] == 'True')

def bn_layer(inputs,is_training,name='BatchNorm',moving_decay=0.9,eps=1e-5):
    shape = inputs.shape
    assert len(shape) in [2,4]

    param_shape = shape[-1]

    gamma = tf.Variable(tf.ones(param_shape), name='gamma')
    beta  = tf.Variable(tf.zeros(param_shape), name='beta')
    mean  = tf.Variable(tf.ones(param_shape), trainable=False, name='mean')
    var   = tf.Variable(tf.ones(param_shape), trainable=False, name='var')

    tf.add_to_collection('l2_losses', tf.contrib.layers.l2_regularizer(lambda1)(gamma)) 
    tf.add_to_collection('l2_losses', tf.contrib.layers.l2_regularizer(lambda1)(beta)) 
    tf.add_to_collection('l2_losses', tf.contrib.layers.l2_regularizer(lambda1)(mean)) 
    tf.add_to_collection('l2_losses', tf.contrib.layers.l2_regularizer(lambda1)(var)) 



    if is_training == True:
        batch_mean, batch_var = tf.nn.moments(inputs,[0,1,2],name='moments')
        mean = tf.assign(mean, batch_mean)
        var = tf.assign(var, batch_var)
        return tf.nn.batch_normalization(inputs,batch_mean+mean*1e-10,batch_var+var*1e-10,gamma,beta,eps)
    else:
        return tf.nn.batch_normalization(inputs,mean,var,gamma,beta,eps)

def get_idx_map(shape):
    """Get index map for a image.
    Args:
        shape: [B, T, H, W] or [B, H, W]
    Returns:
        idx: [B, T, H, W, 2], or [B, H, W, 2]
    """
    s = shape
    ndims = tf.shape(s)
    wdim = ndims - 1
    hdim = ndims - 2
    idx_shape = tf.concat(0, [s, tf.constant([1])])
    ones_h = tf.ones(hdim - 1, dtype='int32')
    ones_w = tf.ones(wdim - 1, dtype='int32')
    h_shape = tf.concat(0, [ones_h, tf.constant([-1]), tf.constant([1, 1])])
    w_shape = tf.concat(0, [ones_w, tf.constant([-1]), tf.constant([1])])

    idx_y = tf.zeros(idx_shape, dtype='float')
    idx_x = tf.zeros(idx_shape, dtype='float')

    h = tf.slice(s, ndims - 2, [1])
    w = tf.slice(s, ndims - 1, [1])
    idx_y += tf.reshape(tf.to_float(tf.range(h[0])), h_shape)
    idx_x += tf.reshape(tf.to_float(tf.range(w[0])), w_shape)
    idx = tf.concat(ndims[0], [idx_y, idx_x])

    return idx

  def _make_evaluation_dict(self, resized_groundtruth_masks=False):
    input_data_fields = fields.InputDataFields
    detection_fields = fields.DetectionResultFields

    image = tf.zeros(shape=[1, 20, 20, 3], dtype=tf.uint8)
    key = tf.constant('image1')
    detection_boxes = tf.constant([[[0., 0., 1., 1.]]])
    detection_scores = tf.constant([[0.8]])
    detection_classes = tf.constant([[0]])
    detection_masks = tf.ones(shape=[1, 1, 20, 20], dtype=tf.float32)
    num_detections = tf.constant([1])
    groundtruth_boxes = tf.constant([[0., 0., 1., 1.]])
    groundtruth_classes = tf.constant([1])
    groundtruth_instance_masks = tf.ones(shape=[1, 20, 20], dtype=tf.uint8)
    if resized_groundtruth_masks:
      groundtruth_instance_masks = tf.ones(shape=[1, 10, 10], dtype=tf.uint8)
    detections = {
        detection_fields.detection_boxes: detection_boxes,
        detection_fields.detection_scores: detection_scores,
        detection_fields.detection_classes: detection_classes,
        detection_fields.detection_masks: detection_masks,
        detection_fields.num_detections: num_detections
    }
    groundtruth = {
        input_data_fields.groundtruth_boxes: groundtruth_boxes,
        input_data_fields.groundtruth_classes: groundtruth_classes,
        input_data_fields.groundtruth_instance_masks: groundtruth_instance_masks
    }
    return eval_util.result_dict_for_single_example(image, key, detections,
                                                    groundtruth)

def test_sample_mvn(session_tf, cov_structure, num_samples):
    """
    Draws 10,000 samples from a distribution
    with known mean and covariance. The test checks
    if the mean and covariance of the samples is
    close to the true mean and covariance.
    """

    N, D = 10000, 2
    means = tf.ones((N, D), dtype=float_type)
    if cov_structure == "full":
        covs = tf.eye(D, batch_shape=[N], dtype=float_type)
    elif cov_structure == "diag":
        covs = tf.ones((N, D), dtype=float_type)

    samples = _sample_mvn(means, covs, cov_structure, num_samples=num_samples)
    value = session_tf.run(samples)

    if num_samples is None:
        assert value.shape == (N, D)
    else:
        assert value.shape == (num_samples, N, D)
        value = value.reshape(-1, D)

    samples_mean = np.mean(value, axis=0)
    samples_cov = np.cov(value, rowvar=False)
    np.testing.assert_array_almost_equal(samples_mean, [1., 1.], decimal=1)
    np.testing.assert_array_almost_equal(samples_cov, [[1., 0.], [0., 1.]], decimal=1)

def _test_logpdf_scalar(scalar):
    x = tf.constant(scalar)
    val_true = stats.norm.logpdf(scalar)
    _assert_eq(norm.logpdf(x), val_true)
    _assert_eq(norm.logpdf(x, tf.zeros([1]), tf.constant(1.0)), val_true)
    _assert_eq(norm.logpdf(x, tf.zeros([1]), tf.ones([1])), val_true)
    _assert_eq(norm.logpdf(x, tf.zeros([1]), tf.diag(tf.ones([1]))), val_true)

 def testDtype(self):
   with self.test_session():
     d = tf.fill([2, 3], 12., name="fill")
     self.assertEqual(d.get_shape(), [2, 3])
     # Test default type for both constant size and dynamic size
     z = tf.ones([2, 3])
     self.assertEqual(z.dtype, tf.float32)
     self.assertEqual([2, 3], z.get_shape())
     self.assertAllEqual(z.eval(), np.ones([2, 3]))
     z = tf.ones(tf.shape(d))
     self.assertEqual(z.dtype, tf.float32)
     self.assertEqual([2, 3], z.get_shape())
     self.assertAllEqual(z.eval(), np.ones([2, 3]))
     # Test explicit type control
     for dtype in (tf.float32, tf.float64, tf.int32,
                   tf.uint8, tf.int16, tf.int8,
                   tf.complex64, tf.complex128, tf.int64, tf.bool):
       z = tf.ones([2, 3], dtype=dtype)
       self.assertEqual(z.dtype, dtype)
       self.assertEqual([2, 3], z.get_shape())
       self.assertAllEqual(z.eval(), np.ones([2, 3]))
       z = tf.ones(tf.shape(d), dtype=dtype)
       self.assertEqual(z.dtype, dtype)
       self.assertEqual([2, 3], z.get_shape())
       self.assertAllEqual(z.eval(), np.ones([2, 3]))

    def __init__(self, num_layers, num_units, batch_size, input_size, keep_prob=1.0):
        self.num_layers = num_layers
        self.grus = []
        self.inits = []
        self.dropout_mask = []
        for layer in range(num_layers):
            input_size_ = input_size if layer == 0 else 2 * num_units
            gru_fw = tf.nn.rnn_cell.MultiRNNCell([
                tf.contrib.cudnn_rnn.CudnnCompatibleGRUCell(num_units=num_units)])

            gru_bw = tf.nn.rnn_cell.MultiRNNCell([
                tf.contrib.cudnn_rnn.CudnnCompatibleGRUCell(num_units=num_units)])

            init_fw = tf.Variable(tf.zeros([num_units]))
            init_fw = tf.expand_dims(tf.tile(tf.expand_dims(init_fw, axis=0), [batch_size, 1]), axis=0)
            init_bw = tf.Variable(tf.zeros([num_units]))
            init_bw = tf.expand_dims(tf.tile(tf.expand_dims(init_bw, axis=0), [batch_size, 1]), axis=0)

            mask_fw = tf.nn.dropout(tf.ones([1, batch_size, input_size_], dtype=tf.float32),
                                    keep_prob=keep_prob)
            mask_bw = tf.nn.dropout(tf.ones([1, batch_size, input_size_], dtype=tf.float32),
                                    keep_prob=keep_prob)

            self.grus.append((gru_fw, gru_bw,))
            self.inits.append((init_fw, init_bw,))
            self.dropout_mask.append((mask_fw, mask_bw,))

  def benchmarkCudnnLSTMTraining(self):
    test_configs = self._GetTestConfig()
    for config_name, config in test_configs.items():
      config = test_configs[config_name]
      num_layers = config["num_layers"]
      num_units = config["num_units"]
      batch_size = config["batch_size"]
      seq_length = config["seq_length"]

      with tf.Graph().as_default(), tf.device("/gpu:0"):
        model = tf.contrib.cudnn_rnn.CudnnLSTM(num_layers, num_units, num_units)
        params_size_t = model.params_size()
        input_data = tf.Variable(tf.ones([seq_length, batch_size, num_units]))
        input_h = tf.Variable(tf.ones([num_layers, batch_size, num_units]))
        input_c = tf.Variable(tf.ones([num_layers, batch_size, num_units]))
        params = tf.Variable(tf.ones([params_size_t]), validate_shape=False)
        output, output_h, output_c = model(
            is_training=True,
            input_data=input_data,
            input_h=input_h,
            input_c=input_c,
            params=params)
        all_grads = tf.gradients([output, output_h, output_c],
                                 [params, input_data, input_h, input_c])
        training_op = tf.group(*all_grads)
        self._BenchmarkOp(training_op, "cudnn_lstm %s %s" %
                          (config_name, self._GetConfigDesc(config)))

    def test_horovod_broadcast(self):
        """Test that the broadcast correctly broadcasts 1D, 2D, 3D tensors."""
        hvd.init()
        rank = hvd.rank()
        size = hvd.size()

        # This test does not apply if there is only one worker.
        if size == 1:
            return

        with self.test_session() as session:
            dtypes = [tf.uint8, tf.int8, tf.uint16, tf.int16,
                      tf.int32, tf.int64, tf.float32, tf.float64,
                      tf.bool]
            dims = [1, 2, 3]
            root_ranks = list(range(size))
            for dtype, dim, root_rank in itertools.product(dtypes, dims, root_ranks):
                try:
                    tensor = tf.ones([17] * dim) * rank
                    root_tensor = tf.ones([17] * dim) * root_rank
                    if dtype == tf.bool:
                        tensor = tensor % 2
                        root_tensor = root_tensor % 2
                    tensor = tf.cast(tensor, dtype=dtype)
                    root_tensor = tf.cast(root_tensor, dtype=dtype)
                    broadcasted_tensor = hvd.broadcast(tensor, root_rank)
                    self.assertTrue(
                        session.run(tf.reduce_all(tf.equal(
                            tf.cast(root_tensor, tf.int32), tf.cast(broadcasted_tensor, tf.int32)))),
                        "hvd.broadcast produces incorrect broadcasted tensor")
                except Exception:
                    import traceback
                    traceback.print_exc()

    def _tf_loss(self, sim, sim_emb):
        """Define loss"""

        if self.use_max_sim_neg:
            max_sim_neg = tf.reduce_max(sim[:, 1:], -1)
            loss = tf.reduce_mean(tf.maximum(0., self.mu_pos - sim[:, 0]) +
                                  tf.maximum(0., self.mu_neg + max_sim_neg))
        else:
            # create an array for mu
            mu = self.mu_neg * np.ones(self.num_neg + 1)
            mu[0] = self.mu_pos

            factors = tf.concat([-1 * tf.ones([1, 1]),
                                 tf.ones([1, tf.shape(sim)[1] - 1])], 1)
            max_margin = tf.maximum(0., mu + factors * sim)
            loss = tf.reduce_mean(tf.reduce_sum(max_margin, -1))

        max_sim_emb = tf.maximum(0., tf.reduce_max(sim_emb, -1))

        loss = (loss +
                # penalize max similarity between intent embeddings
                tf.reduce_mean(max_sim_emb) * self.C_emb +
                # add regularization losses
                tf.losses.get_regularization_loss())
        return loss

        def c_body(c, pa):
          # Zeroing predictions below threshold
          with tf.variable_scope('bboxes_c_select', reuse=True):
            c_scores = b_scores[:, c]
            c_fmask = tf.cast(tf.greater(c_scores, confidence_threshold), scores.dtype)
            c_scores = c_scores * c_fmask
            c_bboxes = b_bboxes * tf.expand_dims(c_fmask, axis=-1)

          # Apply NMS
          with tf.variable_scope('bboxes_c_nms', reuse=True):
            c_indices = tf.image.non_max_suppression(c_bboxes, c_scores, top_k, nms_threshold)
            size = tf.size(c_indices)
            c_batch_ = tf.to_float(b) * tf.ones(shape=[top_k, 1], dtype=tf.float32)  # len(indices) x 1
            c_labels = tf.to_float(c) * tf.ones(shape=[top_k, 1], dtype=tf.float32)  # len(indices) x 1

            extra_size = top_k - size
            c_scores = tf.expand_dims(tf.gather(c_scores, c_indices), axis=-1)  # len(indices) x 1
            empty_c_scores = tf.zeros([extra_size, 1], dtype=tf.float32)
            c_scores = tf.concat([c_scores, empty_c_scores], axis=0)

            c_bboxes = tf.gather(c_bboxes, c_indices)  # len(indices) x 4
            empty_c_bboxes = tf.zeros([extra_size, 4], dtype=tf.float32)
            c_bboxes = tf.concat([c_bboxes, empty_c_bboxes], axis=0)
            c_predictions = tf.concat([c_batch_, c_labels, c_scores, c_bboxes], axis=1)  # len(indices) x 7
          return c + 1, pa.write(index=c - 1, value=c_predictions)

    def __call__(self, inputs, states, scope=None):
        with tf.variable_scope(
                scope or type(self).__name__,
                initializer=tf.random_normal_initializer(stddev=0.01)):
            # get the tensor
            if self._separate_pad:
                t_shape = [self._num_outputs,
                           self._num_outputs,
                           self._num_inputs]
                vec_a = inputs
                vec_b = states
            else:
                t_shape = [self._num_outputs+1,
                           self._num_outputs,
                           self._num_inputs+1]
                vec_a = tf.concat(
                    axis=1, values=[inputs, tf.ones([inputs.get_shape()[0].value, 1])])
                vec_b = tf.concat(
                    axis=1, values=[inputs, tf.ones([inputs.get_shape()[0].value, 1])])
            tensor = get_tt_3_tensor(t_shape, self._ranks, name='W')
            result = bilinear_product_tt_3(vec_a, tensor, vec_b)
            if self._separate_pad:
                # TODO possible weightnorm
                D = tf.get_variable('D', [self._num_inputs, self._num_outputs],
                                    initializer=tf.uniform_unit_scaling_initializer(1.2))
                E = tf.get_variable('E', [self._num_outputs, self._num_outputs],
                                    initializer=tf.uniform_unit_scaling_initializer(1.2))
                b = tf.get_variable('b', [self._num_outputs],
                                    initializer=tf.constant_initializer(0.0))
                z = tf.nn.bias_add(tf.matmul(inputs, D) + tf.matmul(states, E), b)
                result = result + z

            result = self._nonlin(result)
            return result, result

 def decoder_fn(time, cell_state, cell_input, cell_output, context_state):
     with tf.name_scope(name, "simple_decoder_fn_inference",
                        [time, cell_state, cell_input, cell_output,
                         context_state]):
         if cell_input is not None:
             raise ValueError("Expected cell_input to be None, but saw: %s" %
                              cell_input)
         if cell_output is None:
             # invariant that this is time == 0
             next_input_id = tf.ones([batch_size], dtype=dtype) * (
                 start_of_sequence_id)
             done = tf.zeros([batch_size], dtype=tf.bool)
             cell_state = encoder_state
             cell_output = tf.zeros([cell_size],
                                    dtype=tf.float32)
         else:
             softmax_output = output_fn(cell_output)
             if sample:
                 next_input_id = tf.squeeze(tf.multinomial(softmax_output, 1), 1)
             else:
                 next_input_id = tf.argmax(softmax_output, 1)
             next_input_id = tf.cast(next_input_id, dtype=dtype)
             done = tf.equal(next_input_id, end_of_sequence_id)
         next_input = tf.gather(embeddings, next_input_id)
         # if time > maxlen, return all true vector
         done = tf.cond(
             tf.greater(time, maximum_length),
             lambda: tf.ones([batch_size], dtype=tf.bool),
             lambda: done)
         return (done, cell_state, next_input, next_input_id, context_state)

  def _testParamShapes(self, desired_shape):
    tn_param_shapes = tfd.TruncatedNormal.param_shapes(desired_shape)
    # Check the shapes by comparison with the untruncated Normal.
    n_param_shapes = tfd.Normal.param_shapes(desired_shape)
    self.assertAllEqual(
        self.evaluate(tn_param_shapes["loc"]),
        self.evaluate(n_param_shapes["loc"]))
    self.assertAllEqual(
        self.evaluate(tn_param_shapes["scale"]),
        self.evaluate(n_param_shapes["scale"]))
    self.assertAllEqual(
        self.evaluate(tn_param_shapes["low"]),
        self.evaluate(n_param_shapes["loc"]))
    self.assertAllEqual(
        self.evaluate(tn_param_shapes["high"]),
        self.evaluate(n_param_shapes["loc"]))

    loc = tf.zeros(tn_param_shapes["loc"])
    scale = tf.ones(tn_param_shapes["scale"])
    high = tf.ones(tn_param_shapes["high"])
    low = tf.zeros(tn_param_shapes["low"])
    sample_shape = self.evaluate(
        tf.shape(
            tfd.TruncatedNormal(loc=loc, scale=scale, low=low,
                                high=high).sample()))
    self.assertAllEqual(desired_shape, sample_shape)

	def build_rmsprop_optimizer(self, learning_rate, rmsprop_decay, rmsprop_constant, gradient_clip, version):

		with tf.name_scope('rmsprop'):
			optimizer = tf.train.GradientDescentOptimizer(learning_rate)

			grads_and_vars = optimizer.compute_gradients(self.loss)
			grads = [gv[0] for gv in grads_and_vars]
			params = [gv[1] for gv in grads_and_vars]

			if gradient_clip > 0:
				grads = tf.clip_by_global_norm(grads, gradient_clip)

			if version == 'rmsprop':
				return optimizer.apply_gradients(zip(grads, params))
			elif version == 'graves_rmsprop':
				square_grads = [tf.square(grad) for grad in grads]

				avg_grads = [tf.Variable(tf.ones(var.get_shape())) for var in params]
				avg_square_grads = [tf.Variable(tf.ones(var.get_shape())) for var in params]

				update_avg_grads = [grad_pair[0].assign((rmsprop_decay * grad_pair[0]) + ((1 - rmsprop_decay) * grad_pair[1])) 
					for grad_pair in zip(avg_grads, grads)]
				update_avg_square_grads = [grad_pair[0].assign((rmsprop_decay * grad_pair[0]) + ((1 - rmsprop_decay) * tf.square(grad_pair[1]))) 
					for grad_pair in zip(avg_square_grads, grads)]
				avg_grad_updates = update_avg_grads + update_avg_square_grads

				rms = [tf.sqrt(avg_grad_pair[1] - tf.square(avg_grad_pair[0]) + rmsprop_constant)
					for avg_grad_pair in zip(avg_grads, avg_square_grads)]


				rms_updates = [grad_rms_pair[0] / grad_rms_pair[1] for grad_rms_pair in zip(grads, rms)]
				train = optimizer.apply_gradients(zip(rms_updates, params))

				return tf.group(train, tf.group(*avg_grad_updates))

  def _build_iid_normal_model(self, num_timesteps, latent_size,
                              observation_size, transition_variance,
                              observation_variance):
    """Build a model whose outputs are IID normal by construction."""

    transition_variance = self._build_placeholder(transition_variance)
    observation_variance = self._build_placeholder(observation_variance)

    # Use orthogonal matrices to project a (potentially
    # high-dimensional) latent space of IID normal variables into a
    # low-dimensional observation that is still IID normal.
    random_orthogonal_matrix = lambda: np.linalg.qr(
        np.random.randn(latent_size, latent_size))[0][:observation_size, :]
    observation_matrix = self._build_placeholder(random_orthogonal_matrix())

    model = tfd.LinearGaussianStateSpaceModel(
        num_timesteps=num_timesteps,
        transition_matrix=self._build_placeholder(
            np.zeros([latent_size, latent_size])),
        transition_noise=tfd.MultivariateNormalDiag(
            scale_diag=tf.sqrt(transition_variance) *
            tf.ones([latent_size], dtype=self.dtype)),
        observation_matrix=observation_matrix,
        observation_noise=tfd.MultivariateNormalDiag(
            scale_diag=tf.sqrt(observation_variance) *
            tf.ones([observation_size], dtype=self.dtype)),
        initial_state_prior=tfd.MultivariateNormalDiag(
            scale_diag=tf.sqrt(transition_variance) *
            tf.ones([latent_size], dtype=self.dtype)),
        validate_args=True)

    return model

 def p_zt(self, prev_state, t):
   """Computes the model p(z_t| z_{t-1})."""
   batch_size = tf.shape(prev_state)[0]
   if t > 0:
     z_mu_p = prev_state + self.bs[t - 1]
     p_zt = tf.contrib.distributions.Normal(
         loc=z_mu_p, scale=tf.sqrt(tf.ones_like(z_mu_p) * self.variance))
     return p_zt
   else:  # p(z_0) is mixture of two Normals
     mu_pos = tf.ones([batch_size, self.state_size], dtype=self.dtype) * self.prior_mode_mean
     mu_neg = tf.ones([batch_size, self.state_size], dtype=self.dtype) * -self.prior_mode_mean
     z0_pos = tf.contrib.distributions.Normal(
         loc=mu_pos,
         scale=tf.sqrt(tf.ones_like(mu_pos) * self.variance))
     z0_neg = tf.contrib.distributions.Normal(
         loc=mu_neg,
         scale=tf.sqrt(tf.ones_like(mu_neg) * self.variance))
     mode_probs = tf.convert_to_tensor([self.mixing_coeff, 1-self.mixing_coeff], dtype=tf.float64)
     mode_probs = tf.tile(mode_probs[tf.newaxis, tf.newaxis, :], [batch_size, 1, 1])
     mode_selection_dist = tf.contrib.distributions.Categorical(probs=mode_probs)
     z0_dist = tf.contrib.distributions.Mixture(
         cat=mode_selection_dist,
         components=[z0_pos, z0_neg],
         validate_args=False)
     return z0_dist

    def loss_layer(self, project_logits, lengths, name=None):

        with tf.variable_scope("crf_loss" if not name else name):
            small = -1000.0
            start_logits = tf.concat(
                [small * tf.ones(shape=[self.batch_size, 1, self.num_tags]), tf.zeros(shape=[self.batch_size, 1, 1])],
                axis=-1)

            pad_logits = tf.cast(small * tf.ones([self.batch_size, self.num_steps, 1]), tf.float32)
            logits = tf.concat([project_logits, pad_logits], axis=-1)
            logits = tf.concat([start_logits, logits], axis=1)
            targets = tf.concat(
                [tf.cast(self.num_tags * tf.ones([self.batch_size, 1]), tf.int32), self.targets], axis=-1)

            self.trans = tf.get_variable(
                "transitions",
                shape=[self.num_tags + 1, self.num_tags + 1],
                initializer=self.initializer)

            log_likelihood, self.trans = crf_log_likelihood(
                inputs=logits,
                tag_indices=targets,
                transition_params=self.trans,
                sequence_lengths=lengths + 1)

            return tf.reduce_mean(-log_likelihood)

    def getStatsEigen(self, stats=None):
        if len(self.stats_eigen) == 0:
            stats_eigen = {}
            if stats is None:
                stats = self.stats

            tmpEigenCache = {}
            with tf.device('/cpu:0'):
                for var in stats:
                    for key in ['fprop_concat_stats', 'bprop_concat_stats']:
                        for stats_var in stats[var][key]:
                            if stats_var not in tmpEigenCache:
                                stats_dim = stats_var.get_shape()[1].value
                                e = tf.Variable(tf.ones(
                                    [stats_dim]), name='KFAC_FAC/' + stats_var.name.split(':')[0] + '/e', trainable=False)
                                Q = tf.Variable(tf.diag(tf.ones(
                                    [stats_dim])), name='KFAC_FAC/' + stats_var.name.split(':')[0] + '/Q', trainable=False)
                                stats_eigen[stats_var] = {'e': e, 'Q': Q}
                                tmpEigenCache[
                                    stats_var] = stats_eigen[stats_var]
                            else:
                                stats_eigen[stats_var] = tmpEigenCache[
                                    stats_var]
            self.stats_eigen = stats_eigen
        return self.stats_eigen