// enable/disable JIT compilation at different stages.
void EnableXlaCompilation();

// Returns `false` unless `DisableXlaCompilation` was called.
bool FailOnXlaCompilation();

}  // namespace tensorflow

}

static std::atomic<bool> xla_compilation_disabled(false);

void DisableXlaCompilation() { xla_compilation_disabled = true; }

void EnableXlaCompilation() { xla_compilation_disabled = false; }

bool FailOnXlaCompilation() { return xla_compilation_disabled; }

  // round. This is to avoid non-determanism of when compilation is disallowed,
  // for example by changing the threshold.
  if (state == DeviceCompileState::kUncompiled && FailOnXlaCompilation()) {
    VLOG(1) << "XLA compilation disabled: " << function.name() << "\n"
            << absl::StrJoin(
                   args, "\n",
                   [](std::string* out, const XlaCompiler::Argument& arg) {