Home
last modified time | relevance | path

Searched refs:PyBuffer (Results 1 – 12 of 12) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/python/
Dxla.cc269 py::class_<PyBuffer, DeviceArrayBase, std::unique_ptr<PyBuffer>> buffer( in PYBIND11_MODULE()
277 .def_property("_device", &PyBuffer::GetStickyDevice, in PYBIND11_MODULE()
278 &PyBuffer::SetStickyDevice) in PYBIND11_MODULE()
279 .def_property("aval", &PyBuffer::GetAval, &PyBuffer::SetAval) in PYBIND11_MODULE()
286 [](const PyBuffer& pybuffer) -> pybind11::tuple { in PYBIND11_MODULE()
292 [](const PyBuffer& buffer) { in PYBIND11_MODULE()
297 .def_property_readonly("size", &PyBuffer::size) in PYBIND11_MODULE()
298 .def_property_readonly("ndim", &PyBuffer::ndim) in PYBIND11_MODULE()
303 PyBuffer* buffer = buffer_obj.cast<PyBuffer*>(); in PYBIND11_MODULE()
306 .def("copy_to_device", &PyBuffer::CopyToDevice) in PYBIND11_MODULE()
[all …]
Dpy_executable.cc73 StatusOr<std::vector<std::unique_ptr<PyBuffer>>> PyExecutable::PjRtExecute( in PjRtExecute()
81 std::vector<std::unique_ptr<PyBuffer>> outputs; in PjRtExecute()
85 std::make_unique<PyBuffer>(client_, std::move(buffer), traceback)); in PjRtExecute()
90 StatusOr<std::vector<std::unique_ptr<PyBuffer>>> PyExecutable::Execute( in Execute()
91 absl::Span<PyBuffer* const> args) { in Execute()
97 [](PyBuffer* buf) { return buf->buffer(); }); in Execute()
102 std::vector<std::unique_ptr<PyBuffer>> outputs; in Execute()
106 std::make_unique<PyBuffer>(client_, std::move(buffer), traceback)); in Execute()
111 StatusOr<std::vector<std::vector<std::unique_ptr<PyBuffer>>>>
113 absl::Span<const std::vector<PyBuffer*>> args) { in ExecuteOnLocalDevices()
[all …]
Dpy_buffer.cc30 PyBuffer::PyBuffer(std::shared_ptr<PyClient> client, in PyBuffer() function in xla::PyBuffer
45 PyBuffer::~PyBuffer() { in ~PyBuffer()
58 pybind11::tuple PyBuffer::python_shape() const { in python_shape()
62 pybind11::dtype PyBuffer::python_dtype() const { in python_dtype()
67 ClientAndPtr<PjRtDevice> PyBuffer::device() const { in device()
71 StatusOr<std::unique_ptr<PyBuffer>> PyBuffer::CopyToDevice( in CopyToDevice()
81 return std::make_unique<PyBuffer>(dst_device.client, std::move(out), in CopyToDevice()
85 Status PyBuffer::BlockHostUntilReady() { in BlockHostUntilReady()
91 Status PyBuffer::CopyToHostAsync() { in CopyToHostAsync()
108 StatusOr<pybind11::object> PyBuffer::AsNumPyArray(py::handle this_obj) { in AsNumPyArray()
[all …]
Dpy_executable.h59 StatusOr<std::vector<std::unique_ptr<PyBuffer>>> Execute(
60 absl::Span<PyBuffer* const> args);
63 StatusOr<std::vector<std::unique_ptr<PyBuffer>>> PjRtExecute(
67 StatusOr<std::vector<std::vector<std::unique_ptr<PyBuffer>>>>
68 ExecuteOnLocalDevices(absl::Span<const std::vector<PyBuffer*>> args);
73 StatusOr<std::vector<std::vector<std::unique_ptr<PyBuffer>>>>
74 ExecuteShardedOnLocalDevices(absl::Span<const std::vector<PyBuffer*>> args);
Dpy_buffer.h53 class PyBuffer : public DeviceArrayBase {
55 PyBuffer(std::shared_ptr<PyClient> client, std::unique_ptr<PjRtBuffer> buffer,
57 ~PyBuffer();
68 StatusOr<std::unique_ptr<PyBuffer>> CopyToDevice(
131 PyBuffer* next_;
132 PyBuffer* prev_;
Dpy_client.h31 class PyBuffer; variable
109 std::vector<ClientAndPtr<PyBuffer>> LiveBuffers();
131 StatusOr<std::unique_ptr<PyBuffer>> BufferFromPyval(
141 friend class PyBuffer;
149 PyBuffer* buffers_ = nullptr;
Dpy_client.cc57 std::vector<ClientAndPtr<PyBuffer>> PyClient::LiveBuffers() { in LiveBuffers()
59 std::vector<ClientAndPtr<PyBuffer>> buffers; in LiveBuffers()
60 for (PyBuffer* buffer = buffers_; buffer; buffer = buffer->next_) { in LiveBuffers()
146 StatusOr<std::unique_ptr<PyBuffer>> PyClient::BufferFromPyval( in BufferFromPyval()
154 return std::make_unique<PyBuffer>(shared_from_this(), std::move(buffer), in BufferFromPyval()
273 for (PyBuffer* buffer = buffers_; buffer; buffer = buffer->next_) { in HeapProfile()
Ddlpack.cc252 PyBuffer* buffer = py::cast<PyBuffer*>(py_buffer); in BufferToDLPackManagedTensor()
318 StatusOr<std::unique_ptr<PyBuffer>> DLPackManagedTensorToBuffer( in DLPackManagedTensorToBuffer()
367 return std::make_unique<PyBuffer>(std::move(client), std::move(pjrt_buffer), in DLPackManagedTensorToBuffer()
Ddlpack.h31 StatusOr<std::unique_ptr<PyBuffer>> DLPackManagedTensorToBuffer(
Dpmap_lib.cc315 std::vector<std::vector<xla::PyBuffer*>> arg_buffers; in Call()
318 std::vector<xla::PyBuffer*> buffers; in Call()
321 buffers.push_back(py::cast<xla::PyBuffer*>(buf)); in Call()
326 std::vector<std::vector<std::unique_ptr<xla::PyBuffer>>> outputs = in Call()
Djax_jit.cc439 xla::PyBuffer* buffer = py::cast<xla::PyBuffer*>(h); in ArgSignatureOfValue()
649 xla::PyBuffer* buffer = py::cast<xla::PyBuffer*>(obj.attr("device_buffer")); in HandleDeviceArray()
969 if (py::isinstance<xla::PyBuffer>(arg) || in ConvertArgsToBuffers()
971 xla::PyBuffer* buffer; in ConvertArgsToBuffers()
977 buffer = py::cast<xla::PyBuffer*>(arg.attr("device_buffer")); in ConvertArgsToBuffers()
1180 std::vector<std::unique_ptr<xla::PyBuffer>> outputs = in Call()
1244 auto buffer = std::make_unique<xla::PyBuffer>( in BuildJaxjitSubmodule()
/external/python/cpython2/Misc/NEWS.d/
D2.7.4rc1.rst647 structure with data from an object supporting only the old PyBuffer