Home
last modified time | relevance | path

Searched refs:PyTpuExecutable (Results 1 – 3 of 3) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/
H A Dtpu_client_extension.cc143 -> StatusOr<std::unique_ptr<PyTpuExecutable>> { in PYBIND11_MODULE()
145 return PyTpuExecutable::Compile( in PYBIND11_MODULE()
155 -> StatusOr<std::unique_ptr<PyTpuExecutable>> { in PYBIND11_MODULE()
160 return PyTpuExecutable::CompileMlir( in PYBIND11_MODULE()
233 py::class_<PyTpuExecutable>(m, "TpuExecutable") in PYBIND11_MODULE()
235 &PyTpuExecutable::local_logical_device_ids) in PYBIND11_MODULE()
236 .def("local_devices", &PyTpuExecutable::local_devices) in PYBIND11_MODULE()
237 .def_property_readonly("client", &PyTpuExecutable::client) in PYBIND11_MODULE()
239 &PyTpuExecutable::SizeOfGeneratedCodeInBytes) in PYBIND11_MODULE()
240 .def("Delete", &PyTpuExecutable::Delete) in PYBIND11_MODULE()
[all …]
H A Dtpu_client.h333 class PyTpuExecutable {
335 static StatusOr<std::unique_ptr<PyTpuExecutable>> Compile(
341 static StatusOr<std::unique_ptr<PyTpuExecutable>> CompileMlir(
346 PyTpuExecutable(
350 virtual ~PyTpuExecutable() { in ~PyTpuExecutable()
356 PyTpuExecutable(const PyTpuExecutable&) = delete;
357 PyTpuExecutable(PyTpuExecutable&&) = delete;
358 PyTpuExecutable& operator=(const PyTpuExecutable&) = delete;
359 PyTpuExecutable& operator=(PyTpuExecutable&&) = delete;
H A Dtpu_client.cc520 PyTpuExecutable::PyTpuExecutable( in PyTpuExecutable() function in xla::PyTpuExecutable
555 PyTpuExecutable::ExecuteResult PyTpuExecutable::ExecuteHelper( in ExecuteHelper()
632 StatusOr<std::vector<std::unique_ptr<PyTpuBuffer>>> PyTpuExecutable::Execute( in Execute()
677 PyTpuExecutable::ExecuteOnLocalDevices( in ExecuteOnLocalDevices()
765 PyTpuExecutable::ExecuteShardedOnLocalDevices( in ExecuteShardedOnLocalDevices()
804 /*static*/ StatusOr<std::unique_ptr<PyTpuExecutable>> PyTpuExecutable::Compile( in Compile()
872 return std::make_unique<PyTpuExecutable>( in Compile()
877 /*static*/ StatusOr<std::unique_ptr<PyTpuExecutable>>
878 PyTpuExecutable::CompileMlir(mlir::ModuleOp module, in CompileMlir()