input->scalar<ResourceHandle>()() = handle;
    tensors_.push_back(input);
    inputs_.push_back({nullptr, input});
  }

 protected:
  DeviceContext* device_context_;
  Allocator* host_allocator_;
  Allocator* device_allocator_;

  XlaCompiler::Options compiler_options_;
  xla::PjRtClient* pjrt_client_;
  PjRtDeviceCompiler* device_compiler_;

          << inputs.size() << " args";
  std::vector<XlaCompiler::Argument> out;
  out.resize(inputs.size());

  // TODO(cheshire): Avoid duplication with framework/op_kernel.h
  DeviceContext* device_context = nullptr;
  if (device != nullptr) {
    TF_RETURN_IF_ERROR(device->TryGetDeviceContext(&device_context));
    bool using_default_context = false;
    auto cleanup = absl::MakeCleanup([&] {