Searched refs:tpu_replicate_device_names_mapping (Results 1 – 2 of 2) sorted by relevance
522 const TPUReplicateDeviceNamesMapping& tpu_replicate_device_names_mapping);571 TPUReplicateDeviceNamesMapping* tpu_replicate_device_names_mapping,
4167 const TPUReplicateDeviceNamesMapping& tpu_replicate_device_names_mapping) { in LowerOutsideCompilationFunctionalNodes() argument4298 auto iter = tpu_replicate_device_names_mapping.find(replicate); in LowerOutsideCompilationFunctionalNodes()4299 TF_RET_CHECK(iter != tpu_replicate_device_names_mapping.end()); in LowerOutsideCompilationFunctionalNodes()4718 TPUReplicateDeviceNamesMapping* tpu_replicate_device_names_mapping, in RewriteTPUReplicateNode() argument4743 tpu_replicate_device_names_mapping->emplace(replicate, tf_device_assignment); in RewriteTPUReplicateNode()5004 TPUReplicateDeviceNamesMapping tpu_replicate_device_names_mapping; in InternalRun() local5012 &tpu_replicate_device_names_mapping, autotuner_thresh)); in InternalRun()5038 graph, *options.flib_def, tpu_replicate_device_names_mapping)); in InternalRun()