Searched refs:params_info (Results 1 – 2 of 2) sorted by relevance
/external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
D | distributed_tpu_rewrite_pass.cc | 1694 const ParameterInfo& params_info, std::vector<InferredShape>* arg_shapes, in GetArgAndRetvalShapes() argument 1703 params_info.NumInputsToEachReplica()); in GetArgAndRetvalShapes() 1705 arg_shapes->resize(params_info.NumInputsToEachReplica()); in GetArgAndRetvalShapes() 1706 TF_RET_CHECK(input_edges.size() == params_info.NumInputsFromHost()); in GetArgAndRetvalShapes() 1732 for (int64 i = 0; i < params_info.NumReplicas(); ++i) { in GetArgAndRetvalShapes() 1733 for (int64 j = 0; j < params_info.NumPerReplicaArgs(); ++j) { in GetArgAndRetvalShapes() 1738 for (int64 i = 0; i < params_info.NumDistributedArgs(); ++i) { 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() 1752 for (int64 i = 0; i < params_info.NumBroadcastArgs(); ++i) { in GetArgAndRetvalShapes() [all …]
|
D | distributed_tpu_rewrite_pass.h | 295 const ParameterInfo& params_info, std::vector<InferredShape>* arg_shapes, 305 int num_cores_per_replica, const ParameterInfo& params_info, 341 const ParameterInfo& params_info, 353 uint64 library_fingerprint, const ParameterInfo& params_info, 429 const ParameterInfo& params_info, int num_tasks, 532 DataTypeVector* retval_types, ParameterInfo* params_info);
|