Searched refs:host_memory_allocator (Results 1 – 3 of 3) sorted by relevance
207 std::unique_ptr<tensorflow::Allocator> host_memory_allocator, in PjRtStreamExecutorClient() argument213 host_memory_allocator_(std::move(host_memory_allocator)), in PjRtStreamExecutorClient()696 void* staging_buffer = host_memory_allocator()->AllocateRaw( in BufferFromHostBuffer()703 on_delete_callback = [staging_buffer, host_memory_allocator = in BufferFromHostBuffer()704 host_memory_allocator()]() { in BufferFromHostBuffer()705 host_memory_allocator->DeallocateRaw(staging_buffer); in BufferFromHostBuffer()735 void* ptr = host_memory_allocator()->AllocateRaw( in BufferFromHostBuffer()738 ptr, [host_memory_allocator = host_memory_allocator()](void* ptr) { in BufferFromHostBuffer()739 host_memory_allocator->DeallocateRaw(ptr); in BufferFromHostBuffer()
323 auto host_memory_allocator = in GetGpuClient() local339 std::move(host_memory_allocator), in GetGpuClient()
128 std::unique_ptr<tensorflow::Allocator> host_memory_allocator,221 tensorflow::Allocator* host_memory_allocator() const { in host_memory_allocator() function