TFE_OpReset(op_.get(), operation_name, device_.c_str(), status);
    if (TF_GetCode(status) != TF_OK) return;
  }
  TFE_OpAddAttrs(op_.get(), attributes);
  for (int input_index = 0; input_index < inputs.size(); ++input_index) {
    TFE_OpAddInput(op_.get(), inputs[input_index], status);
    if (TF_GetCode(status) != TF_OK) return;
  }
  std::vector<TFE_TensorHandle*> unwrapped_results(expected_max_outputs);

                         int input_index, bool use_function,
                         AbstractTensorHandle** numerical_grad) {
  vector<AbstractTensorHandle*> theta_inputs(inputs.size());
  for (int i{}; i < inputs.size(); ++i) {
    theta_inputs[i] = inputs[i];
  }

  AbstractTensorHandle* theta =
      theta_inputs[input_index];  // parameter we are grad checking

        %weight = "tf.Const"() {value = dense<1> : tensor<1024x24x24x3xi8>} : () -> tensor<1024x24x24x3xi8>
        %input_i32 = "tf.Cast"(%weight) : (tensor<1024x24x24x3xi8>) -> tensor<1024x24x24x3xi32>
        %output = "tf.Sub"(%input_i32, %zp) : (tensor<1024x24x24x3xi32>, tensor<i32>) -> tensor<1024x24x24x3xi32>
        %cast = "tf.Cast"(%output) : (tensor<1024x24x24x3xi32>) -> tensor<1024x24x24x3xf32>

// This interface forwards the reference from input to the output tensors
// corresponding to the indices provided with `input_index` and `output_index`
TF_CAPI_EXPORT extern void TF_OpKernelContext_ForwardRefInputToRefOutput(
    TF_OpKernelContext* ctx, int32_t input_index, int32_t output_index);

// The API releases the opaque lock handle returned with
// `TF_MaybeLockVariableInputMutexesInOrder` API

                                const_op->getResult(0));
    auto dq = rewriter.create<DQ>(const_op->getLoc(), expressed_type, q);
    op.setOperand(input_index, dq.getResult());
    return success();
  }

  LogicalResult replaceStatsOp(
      SourceOp op, quantfork::StatisticsOp stats_op, int input_index,
      const operator_property::TensorProperty& tensor_property,
      PatternRewriter& rewriter) const {

           << absl::StrJoin(xla_output_description.shape.dim_sizes(), ", ")
           << ')';
    if (xla_output_description.input_index >= 0)
      output << " input_index=" << xla_output_description.input_index;
    if (xla_output_description.is_constant) output << " constant";
    if (xla_output_description.is_tensor_list) output << " tensor_list";
    output << '\n';
  }

// CHECK-NEXT:  }

// CHECK:       // InputMapping {0, 1}
// CHECK-NEXT:  // XlaInputShape f32[2]
// CHECK-NEXT:  // XlaInputShape f32[2]
// CHECK-NEXT:  // XlaOutputShape (f32[2])

# CHECK-NEXT:  }

# CHECK:       // InputMapping {0, 1}
# CHECK-NEXT:  // XlaInputShape f32[2]
# CHECK-NEXT:  // XlaInputShape f32[2]
# CHECK-NEXT:  // XlaOutputShape (f32[2])

    Model model, AbstractContext* ctx,
    absl::Span<AbstractTensorHandle* const> inputs, int input_index,
    float* expected_grad, int num_grad, bool use_function,
    double abs_error = 1e-2) {
  Status s;
  AbstractTensorHandlePtr numerical_grad;
  {
    AbstractTensorHandle* numerical_grad_raw;
    s = CalcNumericalGrad(ctx, model, inputs, input_index, use_function,
                          &numerical_grad_raw);

    // Use identity op to avoid the weight being constant-folded.
    %identity = "tf.Identity"(%input) : (tensor<*xi8>) -> tensor<*xi8>
    %input_i32 = "tf.Cast"(%identity) : (tensor<*xi8>) -> tensor<*xi32>
    %output = "tf.Sub"(%input_i32, %zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
    %cast = "tf.Cast"(%output) : (tensor<*xi32>) -> tensor<*xf32>