Searched refs:host_memory_for_device_manager (Results 1 – 4 of 4) sorted by relevance
103 auto* host_memory_for_device_manager = in ExecuteInternal() local105 if (host_memory_for_device_manager == nullptr) { in ExecuteInternal()122 host_memory_for_device_manager, send_callbacks, recv_callbacks)); in ExecuteInternal()201 auto* host_memory_for_device_manager = in ExecuteShardedOnLocalDevicesInternal() local203 if (host_memory_for_device_manager == nullptr) { in ExecuteShardedOnLocalDevicesInternal()221 host_memory_for_device_manager, send_callbacks, recv_callbacks)); in ExecuteShardedOnLocalDevicesInternal()
92 PjRtHostMemoryForDeviceManager* host_memory_for_device_manager) in HostCallbackContext() argument94 host_memory_for_device_manager_(host_memory_for_device_manager), in HostCallbackContext()135 PjRtHostMemoryForDeviceManager* host_memory_for_device_manager,
103 PjRtHostMemoryForDeviceManager* host_memory_for_device_manager, in CreateHostCallbackStateAndAppendSendRecvCallbacks() argument107 std::move(host_callback), host_memory_for_device_manager); in CreateHostCallbackStateAndAppendSendRecvCallbacks()
430 host_memory_for_device_manager) in PjRtClient() argument432 std::move(host_memory_for_device_manager)) {} in PjRtClient()