Searched refs:NumReplicas (Results 1 – 2 of 2) sorted by relevance
159 int64 NumReplicas() const { return num_replicas_; } in NumReplicas() function
1732 for (int64 i = 0; i < params_info.NumReplicas(); ++i) { in GetArgAndRetvalShapes()1986 (params_info.NumReplicas() - 1) * params_info.NumPerReplicaArgs(); in AssignArgsAndRetvalsToCores()2267 params_info.NumPerReplicaArgs() * params_info.NumReplicas() + i - in BuildDynamicShapeNodes()2322 proto.set_num_replicas(params_info.NumReplicas()); in BuildCompileNode()2899 TF_RET_CHECK(params_info.NumReplicas() == tpu_device_names.size()); in BuildExecuteNodes()3057 std::vector<std::vector<Node*>> execute_nodes(params_info.NumReplicas()); in BuildExecuteNodes()3094 for (int64 replica = 0; replica < params_info.NumReplicas(); ++replica) { in BuildExecuteNodes()3124 : params_info.NumReplicas() * in BuildExecuteNodes()