Searched defs:infeed_manager (Results 1 – 2 of 2) sorted by relevance
76 gpu::InfeedManager* infeed_manager = gpu::GetOrCreateInfeedManager(); in EnqueueBuffersToInfeed() local106 gpu::InfeedManager* infeed_manager = gpu::GetOrCreateInfeedManager(); in TransferBufferToInfeedInternal() local
939 def _specialize_model(self, input_specs, infeed_manager): argument1159 def _tpu_model_ops_for_input_specs(self, input_specs, infeed_manager): argument