}

TF_DataType TF_InputDatatype(TF_OpKernelContext* ctx, int index) {
  auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx);
  CHECK_GE(index, 0);                     // Crash OK
  CHECK_LT(index, cc_ctx->num_inputs());  // Crash OK
  return static_cast<TF_DataType>(cc_ctx->input_dtype(index));
}

void TF_SetOutput(TF_OpKernelContext* ctx, int i, const TF_Tensor* tensor,

    const XlaCompiler::ResourceUpdate& write =
        compilation_result->resource_updates[i];
    int actual_input_index = write.input_index - missing_ctx_input_prefix;
    CHECK_GE(actual_input_index, 0);
    CHECK_LT(actual_input_index, ctx->num_inputs());
    Var* var = variable_infos[variable_info_lookup[actual_input_index]].var();
    CHECK(var);

    VLOG(2) << "Updating variable #" << i

  {
    mutex_lock l(graph->mu);
    tensorflow::shape_inference::InferenceContext* ic =
        graph->refiner.GetContext(node);
    CHECK(ic != nullptr);                       // Crash OK
    CHECK_LT(output.index, ic->num_outputs());  // Crash OK
    const auto* shapes_and_types =
        ic->output_handle_shapes_and_types(output.index);
    if (shapes_and_types == nullptr) return nullptr;