Searched refs:PyExecutable (Results 1 – 7 of 7) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | py_executable.cc | 25 PyExecutable::PyExecutable(std::shared_ptr<PyClient> client, in PyExecutable() function in xla::PyExecutable
48 PyExecutable::~PyExecutable() { in ~PyExecutable()
61 std::vector<ClientAndPtr<PjRtDevice>> PyExecutable::AddressableDevices() const { in AddressableDevices()
73 StatusOr<std::vector<std::unique_ptr<PyBuffer>>> PyExecutable::PjRtExecute( in PjRtExecute()
90 StatusOr<std::vector<std::unique_ptr<PyBuffer>>> PyExecutable::Execute( in Execute()
112 PyExecutable::ExecuteOnLocalDevices( in ExecuteOnLocalDevices()
140 PyExecutable::ExecuteShardedOnLocalDevices( in ExecuteShardedOnLocalDevices()
185 StatusOr<std::vector<std::shared_ptr<HloModule>>> PyExecutable::HloModules() in HloModules()
|
| D | py_executable.h | 36 class PyExecutable {
38 PyExecutable(std::shared_ptr<PyClient> client,
42 ~PyExecutable();
99 PyExecutable* next_;
100 PyExecutable* prev_;
|
| D | xla.cc | 338 py::class_<PyExecutable, std::shared_ptr<PyExecutable>> executable( in PYBIND11_MODULE()
340 executable.def_property_readonly("client", &PyExecutable::client) in PYBIND11_MODULE()
342 [](PyExecutable* exec) { in PYBIND11_MODULE()
352 .def("local_devices", &PyExecutable::AddressableDevices) in PYBIND11_MODULE()
354 &PyExecutable::SizeOfGeneratedCodeInBytes) in PYBIND11_MODULE()
355 .def("delete", &PyExecutable::Delete) in PYBIND11_MODULE()
356 .def("execute", &PyExecutable::Execute, py::arg("arguments")) in PYBIND11_MODULE()
357 .def("execute_on_local_devices", &PyExecutable::ExecuteOnLocalDevices, in PYBIND11_MODULE()
360 &PyExecutable::ExecuteShardedOnLocalDevices, py::arg("arguments")) in PYBIND11_MODULE()
361 .def("hlo_modules", &PyExecutable::HloModules) in PYBIND11_MODULE()
[all …]
|
| D | py_client.h | 33 class PyExecutable; variable
135 StatusOr<std::shared_ptr<PyExecutable>> Compile(
142 friend class PyExecutable; variable
150 PyExecutable* executables_ = nullptr;
|
| D | py_client.cc | 158 StatusOr<std::shared_ptr<PyExecutable>> PyClient::Compile( in Compile()
170 return std::make_shared<PyExecutable>( in Compile()
280 for (PyExecutable* executable = executables_; executable; in HeapProfile()
|
| D | pmap_lib.cc | 105 std::shared_ptr<xla::PyExecutable> executable;
244 py::cast<std::shared_ptr<xla::PyExecutable>>(pmap_data["xla_executable"]); in AddCacheEntry()
|
| D | jax_jit.cc | 815 std::shared_ptr<xla::PyExecutable> executable;
1073 auto executable = py::cast<std::shared_ptr<xla::PyExecutable>>( in AddCacheEntry()
|