Searched refs:PyExecutable (Results 1 – 7 of 7) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/python/ |
| H A D | py_executable.cc | 46 PyExecutable::PyExecutable(std::shared_ptr<PyClient> client, in PyExecutable() function in xla::PyExecutable
71 PyExecutable::~PyExecutable() { in ~PyExecutable()
84 std::vector<ClientAndPtr<PjRtDevice>> PyExecutable::AddressableDevices() const { in AddressableDevices()
94 PyExecutable::ExecuteInternal( in ExecuteInternal()
176 PyExecutable::ExecuteWithToken(absl::Span<PyBuffer::object const> args, in ExecuteWithToken()
183 StatusOr<std::vector<PyBuffer::object>> PyExecutable::Execute( in Execute()
192 PyExecutable::ExecuteShardedOnLocalDevicesInternal( in ExecuteShardedOnLocalDevicesInternal()
294 PyExecutable::ExecuteShardedOnLocalDevices( in ExecuteShardedOnLocalDevices()
304 PyExecutable::ExecuteShardedOnLocalDevicesWithTokens( in ExecuteShardedOnLocalDevicesWithTokens()
311 StatusOr<std::vector<std::shared_ptr<HloModule>>> PyExecutable::HloModules() in HloModules()
[all …]
|
| H A D | xla.cc | 359 py::class_<PyExecutable, std::shared_ptr<PyExecutable>> executable( in PYBIND11_MODULE()
361 executable.def_property_readonly("client", &PyExecutable::client) in PYBIND11_MODULE()
363 [](PyExecutable* exec) { in PYBIND11_MODULE()
373 .def("local_devices", &PyExecutable::AddressableDevices) in PYBIND11_MODULE()
375 &PyExecutable::SizeOfGeneratedCodeInBytes) in PYBIND11_MODULE()
376 .def("get_compiled_memory_stats", &PyExecutable::GetCompiledMemoryStats) in PYBIND11_MODULE()
377 .def("delete", &PyExecutable::Delete) in PYBIND11_MODULE()
378 .def("execute", &PyExecutable::Execute, py::arg("arguments"), in PYBIND11_MODULE()
382 .def("execute_with_token", &PyExecutable::ExecuteWithToken, in PYBIND11_MODULE()
385 &PyExecutable::ExecuteShardedOnLocalDevices, py::arg("arguments")) in PYBIND11_MODULE()
[all …]
|
| H A D | py_client.h | 35 class PyExecutable; variable
127 std::vector<std::shared_ptr<PyExecutable>> LiveExecutables();
155 StatusOr<std::shared_ptr<PyExecutable>> Compile(
158 StatusOr<std::shared_ptr<PyExecutable>> CompileMlir(
163 const PyExecutable& executable) const;
164 StatusOr<std::shared_ptr<PyExecutable>> DeserializeExecutable(
169 StatusOr<std::shared_ptr<PyExecutable>> DeserializeExecutable( in DeserializeExecutable()
224 friend class PyExecutable; variable
234 PyExecutable* executables_ = nullptr;
|
| H A D | py_executable.h | 71 class PyExecutable : public std::enable_shared_from_this<PyExecutable> {
73 PyExecutable(std::shared_ptr<PyClient> client,
78 ~PyExecutable();
170 PyExecutable* next_;
171 PyExecutable* prev_;
|
| H A D | py_client.cc | 114 std::vector<std::shared_ptr<PyExecutable>> PyClient::LiveExecutables() { in LiveExecutables()
116 std::vector<std::shared_ptr<PyExecutable>> executables; in LiveExecutables()
117 for (PyExecutable* exec = executables_; exec; exec = exec->next_) { in LiveExecutables()
318 StatusOr<std::shared_ptr<PyExecutable>> PyClient::Compile( in Compile()
331 return std::make_shared<PyExecutable>( in Compile()
336 StatusOr<std::shared_ptr<PyExecutable>> PyClient::CompileMlir( in CompileMlir()
352 return std::make_shared<PyExecutable>( in CompileMlir()
358 const PyExecutable& executable) const { in SerializeExecutable()
362 StatusOr<std::shared_ptr<PyExecutable>> PyClient::DeserializeExecutable( in DeserializeExecutable()
375 return std::make_shared<PyExecutable>( in DeserializeExecutable()
[all …]
|
| H A D | pmap_lib.cc | 169 std::shared_ptr<xla::PyExecutable> executable;
364 std::shared_ptr<xla::PyExecutable> executable; in PopulateCacheEntry()
366 executable = py::cast<std::shared_ptr<xla::PyExecutable>>( in PopulateCacheEntry()
|
| H A D | jax_jit.cc | 332 std::shared_ptr<xla::PyExecutable> executable;
781 auto executable = py::cast<std::shared_ptr<xla::PyExecutable>>( in PopulateCacheEntry()
|