Add a method to PyClient returning std::shared_ptr<PjRtClient>.

PiperOrigin-RevId: 318092346
Change-Id: Ic5cdbc20a270f7a2c62eebbd47a2a0591f4d7c46
This commit is contained in:
A. Unique TensorFlower 2020-06-24 10:33:26 -07:00 committed by TensorFlower Gardener
parent 1b049d2a22
commit fcad3004bc

View File

@ -91,6 +91,7 @@ class PyClient : public std::enable_shared_from_this<PyClient> {
explicit PyClient(std::shared_ptr<PjRtClient> pjrt_client);
PjRtClient* pjrt_client() const { return pjrt_client_.get(); }
std::shared_ptr<PjRtClient> shared_pjrt_client() { return pjrt_client_; }
const std::string& platform_name() const {
return pjrt_client_->platform_name();