DCHECK(while_ctx != nullptr);

  // Record 'summed_grads' as the backprop input associated with 'exit_node'
  std::map<Node*, Output>& backprops = while_backprops_[while_ctx];
  DCHECK(backprops.find(exit_node) == backprops.end());
  backprops[exit_node] = summed_grads;

  // Wait until we have all exit nodes' backprops collected before processing
  // the while loop.

  // Build new BackPropFilterOp.
  auto loc = backprop.getLoc();
  auto new_backprop = builder.create<TF::Conv2DBackpropFilterOp>(
      loc, new_result_type, input, new_filter_sizes, backprop.getOutBackprop(),
      strides, backprop.getUseCudnnOnGpu(), backprop.getPadding(),
      backprop.getExplicitPaddings(), backprop.getDataFormat(),
      backprop.getDilations());

// computes loss and backprop of the loss with respect to 'features'.
//
// Softmax cross entropy loss is defined as follows:
//
//  loss = Sum(-labels * Log(Exp(features) / Sum(Exp(features)))
//  loss = Sum(-labels * LogSoftmax(features))
//
// Computing gradient of the loss with respect to features gives us,
//
//  backprop = (Exp(features) / Sum(Exp(features))) - labels

// any gradients to be computed).
//
// Finally, we start a backprop stack with a set of tape entries for which we
// have all gradients available. This set usually is a subset of the set of
// targets (not all since targets which have outputs in the tape will not have
// gradients available initially).
//
// Then we repeatedly pop an entry from the stack, run its backprop, and update

    // CHECK: %[[offset_backprop:.*]] = mhlo.convert %[[red2]] : tensor<8xf32>

    // CHECK: %[[x_backprop:.*]] = mhlo.convert %[[mul3]] : tensor<8x8x8x8xf32>
    // CHECK: return %[[x_backprop]] : tensor<8x8x8x8xf32>

      Add(params_->body_inputs[0], {one, 0}, params_->body_graph, s_);
  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);
  params_->body_outputs[0] = {add, 0};

  ExpectOK();

  // Create backprop graph
  TF_Output grad_output;
  TF_AddGradients(graph_, outputs_.data(), outputs_.size(), inputs_.data(), 1,
                  nullptr, s_, &grad_output);
  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);