/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/batching_util/ |
H A D | serial_device_batch_scheduler_test.cc | 101 auto queue_callback = [&mu, &processed_batches, &finish_processing]( in TEST() local 124 TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue1)); in TEST() 125 TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue2)); in TEST() 126 TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue3)); in TEST() 150 auto queue_callback = [&mu, &processed_batches, &start_processing, &pending, in TEST() local 188 TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue)); in TEST() 210 auto queue_callback = in TEST() local 239 TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue1)); in TEST() 241 TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue2)); in TEST() 243 TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue3)); in TEST() [all …]
|
H A D | adaptive_shared_batch_scheduler_test.cc | 107 auto queue_callback = [&mu, &processed_batches, &finish_processing]( in TEST() local 128 TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue)); in TEST() 150 auto queue_callback = [&env](std::unique_ptr<Batch<FakeTask>> batch) { in TEST() local 171 TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue)); in TEST() 217 auto queue_callback = [&mu, &processed_batches, &finish_processing]( in TEST() local 247 TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue1)); in TEST() 249 TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue2)); in TEST() 292 auto queue_callback = [&mu, &processed_batches, &finish_processing]( in TEST() local 322 TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue1)); in TEST() 324 TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue2)); in TEST() [all …]
|
/aosp_15_r20/external/pytorch/torch/distributed/algorithms/ddp_comm_hooks/ |
H A D | mixed_precision_hooks.py | 84 Variable._execution_engine.queue_callback(wait_for_stream_cb)
|
H A D | optimizer_overlap_hooks.py | 112 Variable._execution_engine.queue_callback(wait_for_optim_stream_callback)
|
/aosp_15_r20/external/pytorch/torch/_dynamo/ |
H A D | external_utils.py | 114 def queue_callback(final_callbacks, cb): member in FakeCompiledAutogradEngine
|
/aosp_15_r20/external/pytorch/torch/_functorch/_aot_autograd/ |
H A D | logging_utils.py | 107 … torch.autograd.variable.Variable._execution_engine.queue_callback( # type: ignore[attr-defined]
|
/aosp_15_r20/external/pytorch/torch/utils/ |
H A D | module_tracker.py | 74 torch.autograd.Variable._execution_engine.queue_callback(callback)
|
/aosp_15_r20/external/pytorch/torch/csrc/autograd/ |
H A D | engine.h | 188 void queue_callback(std::function<void()> callback);
|
H A D | python_engine.cpp | 407 engine.queue_callback([callback]() { in THPEngine_queue_callback()
|
H A D | engine.cpp | 1410 void Engine::queue_callback(std::function<void()> callback) { in queue_callback() function in torch::autograd::Engine
|
/aosp_15_r20/external/pytorch/torch/distributed/_tools/ |
H A D | mod_tracker.py | 76 torch.autograd.Variable._execution_engine.queue_callback(callback)
|
/aosp_15_r20/external/pytorch/torch/distributed/_composable/fsdp/ |
H A D | _fsdp_state.py | 329 Variable._execution_engine.queue_callback(
|
/aosp_15_r20/external/pytorch/torch/nn/parallel/ |
H A D | distributed.py | 266 Variable._execution_engine.queue_callback( # type: ignore[call-arg,misc]
|
/aosp_15_r20/external/pytorch/torch/distributed/fsdp/ |
H A D | _runtime_utils.py | 1525 Variable._execution_engine.queue_callback(
|
/aosp_15_r20/external/pytorch/torch/csrc/distributed/c10d/ |
H A D | reducer.cpp | 916 torch::autograd::Engine::get_default_engine().queue_callback([this] { in mark_variable_ready()
|
/aosp_15_r20/external/pytorch/torch/_dynamo/variables/ |
H A D | misc.py | 951 torch._dynamo.external_utils.FakeCompiledAutogradEngine.queue_callback,
|
/aosp_15_r20/external/pytorch/test/inductor/ |
H A D | test_compiled_autograd.py | 2209 torch.autograd.Variable._execution_engine.queue_callback(callback_final)
|
/aosp_15_r20/external/pytorch/test/ |
H A D | test_autograd.py | 7328 Variable._execution_engine.queue_callback(callback_final) 7338 Variable._execution_engine.queue_callback(callback_adds_callback) 7353 torch.autograd.Variable._execution_engine.queue_callback(callback) 7384 Variable._execution_engine.queue_callback(inc_inner_counter) 7398 Variable._execution_engine.queue_callback(inc_outer_counter)
|
H A D | test_cuda.py | 1326 lambda grad: torch.autograd.Variable._execution_engine.queue_callback(
|
/aosp_15_r20/external/pytorch/torch/_C/ |
H A D | __init__.pyi.in | 1675 def queue_callback(self, callback: Callable[[], None]) -> None: ...
|