Create tensorflow op for running one step of descent.

def prepare_for_optimization(self):
    """Create tensorflow op for running one step of descent."""
    if self.params['eig_type'] == 'TF':
      self.eig_vec_estimate = self.get_min_eig_vec_proxy()
    elif self.params['eig_type'] == 'LZS':
      self.eig_vec_estimate = self.dual_object.m_min_vec
    else:
      self.eig_vec_estimate = tf.placeholder(tf.float32, shape=(self.dual_object.matrix_m_dimension, 1))
    self.stopped_eig_vec_estimate = tf.stop_gradient(self.eig_vec_estimate)
    # Eig value is v^\top M v, where v is eigen vector
    self.eig_val_estimate = tf.matmul(
        tf.transpose(self.stopped_eig_vec_estimate),
        self.dual_object.get_psd_product(self.stopped_eig_vec_estimate))
    # Penalizing negative of min eigen value because we want min eig value
    # to be positive
    self.total_objective = (
        self.dual_object.unconstrained_objective
        + 0.5 * tf.square(
            tf.maximum(-self.penalty_placeholder * self.eig_val_estimate, 0)))
    global_step = tf.Variable(0, trainable=False)
    # Set up learning rate as a placeholder
    self.learning_rate = tf.placeholder(tf.float32, shape=[])

    # Set up the optimizer
    if self.params['optimizer'] == 'adam':
      self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
    elif self.params['optimizer'] == 'adagrad':
      self.optimizer = tf.train.AdagradOptimizer(learning_rate=self.learning_rate)
    elif self.params['optimizer'] == 'momentum':
      self.optimizer = tf.train.MomentumOptimizer(
          learning_rate=self.learning_rate,
          momentum=self.params['momentum_parameter'],
          use_nesterov=True)
    else:
      self.optimizer = tf.train.GradientDescentOptimizer(
          learning_rate=self.learning_rate)

    # Write out the projection step
    self.train_step = self.optimizer.minimize(
        self.total_objective, global_step=global_step)

    self.sess.run(tf.global_variables_initializer())

    # Projecting the dual variables
    proj_ops = []
    for i in range(self.dual_object.nn_params.num_hidden_layers + 1):
      # Lambda_pos is non negative for switch indices,
      # Unconstrained for positive indices
      # Zero for negative indices
      proj_ops.append(self.dual_object.lambda_pos[i].assign(
          tf.multiply(self.dual_object.positive_indices[i],
                      self.dual_object.lambda_pos[i])+
          tf.multiply(self.dual_object.switch_indices[i],
                      tf.nn.relu(self.dual_object.lambda_pos[i]))))
      proj_ops.append(self.dual_object.lambda_neg[i].assign(
          tf.multiply(self.dual_object.negative_indices[i],
                      self.dual_object.lambda_neg[i])+
          tf.multiply(self.dual_object.switch_indices[i],
                      tf.nn.relu(self.dual_object.lambda_neg[i]))))
      # Lambda_quad is only non zero and positive for switch
      proj_ops.append(self.dual_object.lambda_quad[i].assign(
          tf.multiply(self.dual_object.switch_indices[i],
                      tf.nn.relu(self.dual_object.lambda_quad[i]))))
      # Lambda_lu is always non negative
      proj_ops.append(self.dual_object.lambda_lu[i].assign(
          tf.nn.relu(self.dual_object.lambda_lu[i])))

    self.proj_step = tf.group(proj_ops)

    # Create folder for saving stats if the folder is not None
    if (self.params.get('stats_folder') and
        not tf.gfile.IsDirectory(self.params['stats_folder'])):
      tf.gfile.MkDir(self.params['stats_folder'])