FunctionType func_type = func.getFunctionType();
    int num_inputs = func_type.getNumInputs();
    int num_results = func_type.getNumResults();

    // For each argument type in function's arguments, change it to uranked
    // tensor type if it's a variant type.
    SmallVector<Type, 8> updated_argument_types;
    updated_argument_types.reserve(num_inputs);
    UpdateTensorListTypes<mlir::OperandRange>(

}

void TFE_OpAddInputList(TFE_Op* op, TFE_TensorHandle** inputs, int num_inputs,
                        TF_Status* status) {
  status->status = tensorflow::unwrap(op)->AddInputList(
      {reinterpret_cast<tensorflow::AbstractTensorHandle**>(
           tensorflow::unwrap(inputs)),
       static_cast<size_t>(num_inputs)});
}

extern int TFE_OpGetFlatInputCount(const TFE_Op* op, TF_Status* status) {

// TODO(ycao): Support computation with compile-time constant, which requires
// non-trivial input mapping as implemented now.
void GetInputMappingForMlir(int num_inputs, std::vector<int>* input_mapping) {
  input_mapping->resize(num_inputs, 0);
  std::iota(input_mapping->begin(), input_mapping->end(), 0);
}

static void RegisterDialects(mlir::DialectRegistry& registry) {

  }
}

bool TF_IsRefInput(TF_OpKernelContext* ctx, int i, TF_Status* status) {
  auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx);
  if (i < 0 || i >= cc_ctx->num_inputs()) {
    TF_SetStatus(status, TF_OUT_OF_RANGE, "input index out of range");
    return false;
  }
  TF_SetStatus(status, TF_OK, "");
  return cc_ctx->input_is_ref(i);
}

#ifndef IS_MOBILE_PLATFORM

    : OpKernel(ctx), platform_info_(XlaPlatformInfoFromDevice(ctx->device())) {}

void XlaRunOp::Compute(OpKernelContext* ctx) {
  VLOG(3) << "XlaRunOp " << def().name();
  Tensor key_tensor = ctx->input(ctx->num_inputs() - 1);

  bool use_pjrt =
      GetXlaOpsCommonFlags()
          ->tf_xla_use_device_api.IsEnabledInXlaCompileAndRunForDevice(
              platform_info_.device_type());

  if (use_pjrt) {

    mlir::FunctionType callee_type = callee.getFunctionType();
    SmallVector<Type, 4> expected_callee_input_types;
    auto num_inputs = op.getInputs().size() / 2;
    for (unsigned i = 0; i < num_inputs; ++i) {
      auto input_type =
          mlir::dyn_cast<RankedTensorType>(op.getOperand(i).getType());
      auto scalar = RankedTensorType::get({}, input_type.getElementType());

                                        ? mhlo::Precision::DEFAULT
                                        : mlir::mhlo::Precision::HIGHEST;
  llvm::SmallVector<mlir::Attribute, 2> attr_vec;
  const int num_inputs = 2;
  for (int i = 0; i < num_inputs; i++) {
    attr_vec.push_back(
        mlir::mhlo::PrecisionAttr::get(builder->getContext(), precision));
  }
  return builder->getArrayAttr(attr_vec);
}

const std::vector<AttrView>& OpView::Attributes() const {
  return argument_attrs_;
}

const std::vector<OpArgumentView>& OpView::AllArguments() const {
  return all_arguments_;
}

int OpView::NumInputs() const { return input_args_.size(); }

int OpView::NumOutputs() const { return output_args_.size(); }

ArgView OpView::OnlyInput() const {
  CHECK_EQ(input_args_.size(), 1);  // Crash OK
  return input_args_.front();

      }
    }

    return mlir::UnrankedTensorType::get(element_type);
  };

  // Below we only try and do some shape inference for "source" ops which have
  // no inputs.
  if (node.num_inputs() > 0) return default_type();

  // Do some simply inference here to get the function arguments correct for
  // this common case.
  // TODO(jpienaar): Reconsider post refactoring shape functions.

      IsSingleCallRegion(while_region.getBody()), while_arg_matcher);

  // All existing inputs to while region are inputs to the functional while.
  auto new_inputs = llvm::to_vector<4>(while_region.getOperands());

  // All existing results will also be generated by the functional while.
  auto new_result_types = llvm::to_vector<4>(while_region.getResultTypes());