return result;
}

// The backprop loop counter and main backprop loop run in their own execution
// frame (conceptually, the main forward loop and forward loop counter run
// together in a frame, then the backprop loop counter and backprop loop run
// together in a different frame). This returns the frame name to use for the
// backprop while loops.

  *loss = temp_outputs[0];
  *backprop = temp_outputs[1];
  return status;
}

// Op: ReluGrad()
// Summary: Computes rectified linear gradients for a Relu operation.
//
// Description:
Status ReluGrad(AbstractContext* ctx, AbstractTensorHandle* const gradients,
                AbstractTensorHandle* const features,
                AbstractTensorHandle** backprops, const char* name,

    AbstractTensorHandle** backprop, const char* name = nullptr,
    const char* raw_device_name = nullptr);

// Computes rectified linear gradients for a Relu operation.
Status ReluGrad(AbstractContext* ctx, AbstractTensorHandle* const gradients,
                AbstractTensorHandle* const features,
                AbstractTensorHandle** backprops, const char* name = nullptr,

    # compute with op.
    with backprop.GradientTape() as gt:
      for var_ in vars_:
        gt.watch(var_)
      y = compute_op(**op_kwargs)  # uses op and decomposites by the graph pass.
      grads = gt.gradient(y, vars_)  # uses registered gradient function.

    # compute with composition
    with backprop.GradientTape() as gt:
      for var_ in vars_:
        gt.watch(var_)

  %arg2: tensor<64xf32>
) -> tensor<1x28x28x64xf32> {

  // CHECK: "tf.FusedBatchNormGradV3"
  // CHECK-SAME: (%[[X_TRANSPOSE:[0-9]*]], %[[Y_TRANSPOSE:[0-9]*]],
  // CHECK-SAME: data_format = "NCHW"
  %x_backprop, %scale_backprop, %offset_backprop, %reserve_1, %reserve_2
    = "tf.FusedBatchNormGradV3"(%arg0, %arg1, %arg2, %arg2, %arg2, %arg2)
       {
         data_format = "NHWC",
         epsilon = 1.001 : f32,

  %arg2: tensor<64xf32>
) -> tensor<1x28x28x64xf32> {

  // CHECK: "tf.FusedBatchNormGradV3"
  // CHECK-SAME: (%[[X_TRANSPOSE:[0-9]*]], %[[Y_TRANSPOSE:[0-9]*]],
  // CHECK-SAME: data_format = "NCHW"
  %x_backprop, %scale_backprop, %offset_backprop, %reserve_1, %reserve_2
    = "tf.FusedBatchNormGradV3"(%arg0, %arg1, %arg2, %arg2, %arg2, %arg2)
       {
         data_format = "NHWC",
         epsilon = 1.001 : f32,

  // CHECK: %[[RES_PERM:.*]] = "tf.Const"()
  // CHECK-SAME: <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}>

  // CHECK: %[[RES_TPOSE:[0-9]*]] = "tf.Transpose"
  // CHECK-SAME: (%x_backprop, %[[RES_PERM]])
  // CHECK: return %[[RES_TPOSE]]

  %x_backprop, %scale_backprop, %offset_backprop, %reserve_1, %reserve_2
    = "tf.FusedBatchNormGradV3"(%arg0, %arg1, %arg2, %arg2, %arg2, %arg2)
       {
         data_format = "NHWC",

    absl::Span<AbstractTensorHandle*> outputs) {
  AbstractTensorHandle* loss;
  AbstractTensorHandle* backprop;
  TF_RETURN_IF_ERROR(ops::SparseSoftmaxCrossEntropyWithLogits(
      ctx, inputs[0], inputs[1], &loss, &backprop,
      "SparseSoftmaxCrossEntropyWithLogits"));
  // `gradient_checker` only works with model that returns only 1 tensor.

//
Status SparseSoftmaxCrossEntropyWithLogits(AbstractContext* ctx, AbstractTensorHandle* const features, AbstractTensorHandle* const labels, AbstractTensorHandle** loss, AbstractTensorHandle** backprop, const char* name = nullptr, const char* raw_device_name = nullptr);

//

// Summary:
//
// Description:
Status SparseSoftmaxCrossEntropyWithLogits(AbstractContext* ctx, AbstractTensorHandle* const features, AbstractTensorHandle* const labels, AbstractTensorHandle** loss, AbstractTensorHandle** backprop, const char* name, const char* raw_device_name) {
  AbstractOperationPtr op_ptr(ctx->CreateOperation());
  TF_RETURN_IF_ERROR(op_ptr->Reset("SparseSoftmaxCrossEntropyWithLogits", raw_device_name));