def _formulate_smt_constraints_final_layer(
    z3_optimizer, smt_output, delta, label_index):
  """Formulates smt constraints using the logits in the final layer.

  Generates constraints by setting the logit corresponding to label_index to be
  more than the rest of the logits by an amount delta for the forward pass of a
  masked image.

  Args:
    z3_optimizer: instance of z3.Optimizer, z3 optimizer.
    smt_output: list of z3.ExprRef with length num_output.
    delta: float, masked logit for the label is greater than the rest of the
        masked logits by delta.
    label_index: int, index of the label of the training image.

  Returns:
    z3 optimizer with added smt constraints.
  """
  for i, output in enumerate(smt_output):
    if i != label_index:
      z3_optimizer.solver.add(smt_output[label_index] - output > delta)
  return z3_optimizer