Searched refs:NumPerReplicaArgs (Results 1 – 2 of 2) sorted by relevance
1733 for (int64 j = 0; j < params_info.NumPerReplicaArgs(); ++j) { in GetArgAndRetvalShapes()1739 TF_RETURN_IF_ERROR(check_shape(params_info.NumPerReplicaArgs() + i)); in GetArgAndRetvalShapes()1743 i < params_info.NumPerReplicaArgs() + params_info.NumDistributedArgs(); in GetArgAndRetvalShapes()1756 (*arg_shapes)[i + params_info.NumPerReplicaArgs() + in GetArgAndRetvalShapes()1769 (*arg_shapes)[i + params_info.NumPerReplicaArgs() + in GetArgAndRetvalShapes()1788 (*arg_shapes)[i + params_info.NumPerReplicaArgs() + in GetArgAndRetvalShapes()1986 (params_info.NumReplicas() - 1) * params_info.NumPerReplicaArgs(); in AssignArgsAndRetvalsToCores()2246 TF_RET_CHECK(num_args == params_info.NumPerReplicaArgs() + in BuildDynamicShapeNodes()2267 params_info.NumPerReplicaArgs() * params_info.NumReplicas() + i - in BuildDynamicShapeNodes()2268 params_info.NumPerReplicaArgs(); in BuildDynamicShapeNodes()[all …]
161 int64 NumPerReplicaArgs() const { return num_per_replica_args_; } in NumPerReplicaArgs() function