Searched refs:tpu_replicate_device_names_mapping (Results 1 – 2 of 2) sorted by relevance
510 const TPUReplicateDeviceNamesMapping& tpu_replicate_device_names_mapping);558 TPUReplicateDeviceNamesMapping* tpu_replicate_device_names_mapping,
3713 const TPUReplicateDeviceNamesMapping& tpu_replicate_device_names_mapping) { in LowerOutsideCompilationFunctionalNodes() argument3844 auto iter = tpu_replicate_device_names_mapping.find(replicate); in LowerOutsideCompilationFunctionalNodes()3845 TF_RET_CHECK(iter != tpu_replicate_device_names_mapping.end()); in LowerOutsideCompilationFunctionalNodes()4217 TPUReplicateDeviceNamesMapping* tpu_replicate_device_names_mapping, in RewriteTPUReplicateNode() argument4241 tpu_replicate_device_names_mapping->emplace(replicate, tf_device_assignment); in RewriteTPUReplicateNode()4478 TPUReplicateDeviceNamesMapping tpu_replicate_device_names_mapping; in Run() local4486 &tpu_replicate_device_names_mapping, autotuner_thresh)); in Run()4512 graph, *options.flib_def, tpu_replicate_device_names_mapping)); in Run()