void runOnOperation() override;
};

constexpr char kAliasingAttr[] = "tf.aliasing_output";
constexpr int kUnassigned = -1;

struct AliasInfo {
  AliasInfo() : input_index(kUnassigned), output_index(kUnassigned) {}
  int input_index;
  int output_index;
};

// Idenitfy tf_device.cluster_func input-output alias pairs.
// This is currently conservative, primarily handling the following base case:
// ```

namespace tensorflow {
namespace gradients {

/* Returns numerical grad inside `dtheta_approx` given `forward` model and
 * parameter specified by `input_index`.
 *
 * I.e. if y = <output of the forward model> and w = inputs[input_index],
 * this will calculate dy/dw numerically.
 *
 * `use_function` indicates whether to use graph mode(true) or eager(false).
 *

                         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

// 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

// 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>

    if (!outputs[i]) continue;

    AbstractTensorHandlePtr numerical_grad;
    {
      AbstractTensorHandle* numerical_grad_raw;
      s = CalcNumericalGrad(ctx, model, inputs,
                            /*input_index=*/i, use_function,
                            &numerical_grad_raw);
      ASSERT_EQ(errors::OK, s.code()) << s.message();
      numerical_grad.reset(numerical_grad_raw);
    }

    TF_Tensor* numerical_tensor;