ASSERT_EQ(errors::OK, status_.code()) << status_.message();

  auto ReluGradModel = BuildGradModel(ReluModel, registry_);

  float X_vals[] = {1.0f, 2.0f, 3.0f, -5.0f, -4.0f, -3.0f, 2.0f, 10.0f, -1.0f};
  int64_t X_dims[] = {3, 3};
  AbstractTensorHandlePtr X;
  {
    AbstractTensorHandle* X_raw;
    status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
        immediate_execution_ctx_.get(), X_vals, X_dims, 2, &X_raw);

                    value_vec[1].AsInt64(), llvm::ArrayRef<int64_t>(vec1),
                    value_vec[3].AsInt64(), value_vec[4].AsInt64(),
                    llvm::ArrayRef<int64_t>(vec2), value_vec[6].AsInt64(),
                    value_vec[7].AsInt64(), llvm::ArrayRef<int64_t>(vec3));
            auto named_attr =
                builder->getNamedAttr(key, conv_dimension_numbers_attr);

      requested_read_length = INT32_MAX;
    else
      requested_read_length = n;

    // `pread` returns a `ssize_t` on POSIX, but due to interface being
    // cross-platform, return type of `Read` is `int64_t`.
    int64_t r = int64_t{pread(posix_file->fd, dst, requested_read_length,
                              static_cast<off_t>(offset))};
    if (r > 0) {
      dst += r;
      offset += static_cast<uint64_t>(r);

}

// Checks whether the stride array matches the layout of compact, row-majored
// data.
bool IsValidStrideCompactRowMajorData(int64_t* shape_arr, int64_t* stride_arr,
                                      int ndim) {
  bool valid = true;
  int64_t expected_stride = 1;
  for (int i = ndim - 1; i >= 0; --i) {
    // Empty tensors are always compact regardless of strides.
    if (shape_arr[i] == 0) return true;

  return num_dims;
}

int64_t TFE_TensorHandleNumElements(TFE_TensorHandle* h, TF_Status* status) {
  if (h == nullptr) {
    status->status = tensorflow::errors::InvalidArgument("Invalid handle");
    return -1;
  }

  int64_t num_elements = -1;
  status->status = tensorflow::unwrap(h)->NumElements(&num_elements);
  return num_elements;
}

int64_t TFE_TensorHandleDim(TFE_TensorHandle* h, int dim_index,

    set_arg_data(1, data);
  }
  int64_t* arg1_data() {
    return static_cast<int64_t*>(arg_data(1));
  }
  int64_t& arg1(size_t dim0, size_t dim1) {
    return (*static_cast<int64_t(*)[3][4]>(
        arg_data(1)))[dim0][dim1];
  }
  const int64_t* arg1_data() const {
    return static_cast<const int64_t*>(arg_data(1));
  }
  const int64_t& arg1(size_t dim0, size_t dim1) const {

};

struct TFE_MonitoringIntGauge0 : TFE_MonitoringGauge<int64_t, 0> {
  using TFE_MonitoringGauge::TFE_MonitoringGauge;
};
struct TFE_MonitoringIntGauge1 : TFE_MonitoringGauge<int64_t, 1> {
  using TFE_MonitoringGauge::TFE_MonitoringGauge;
};
struct TFE_MonitoringIntGauge2 : TFE_MonitoringGauge<int64_t, 2> {
  using TFE_MonitoringGauge::TFE_MonitoringGauge;
};

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"

namespace mlir::quant {
namespace {

void PackOperandTestHelper(
    const llvm::SmallVector<int64_t>& unpacked_shape,
    const llvm::SmallVector<int8_t>& unpacked_values, int pack_dim,
    const llvm::SmallVector<int64_t>& expected_packed_shape,
    const llvm::SmallVector<int8_t>& expected_packed_values) {
  MLIRContext context;

  auto* cc_agent = reinterpret_cast<tsl::CoordinationServiceAgent*>(agent);
  return cc_agent->IsInitialized();
}

void TF_CoordinationServiceInsertKeyValue(const char* key, int64_t key_size,
                                          const char* value, int64_t value_size,
                                          TF_CoordinationServiceAgent* agent,
                                          TF_Status* status) {

llvm::SmallVector<std::pair<int64_t, llvm::SmallVector<Value, 4>>, 4>
AnnotateCompileOpAndGetExecuteArgToWhileArgsMapping(
    TF::WhileRegionOp while_op, tf_device::ReplicateOp replicate,
    TF::TPUExecuteAndUpdateVariablesOp execute,
    tf_device::LaunchOp compile_launch) {
  Region& body = while_op.getBody();
  Region& cond = while_op.getCond();

  llvm::SmallVector<std::pair<int64_t, llvm::SmallVector<Value, 4>>, 4> mapping;