Searched refs:retval_shardings (Results 1 – 3 of 3) sorted by relevance
101 std::map<int, xla::OpSharding> retval_shardings; in ComputeArgAndRetvalShardings() local116 retval_shardings[index] = std::move(*sharding); in ComputeArgAndRetvalShardings()119 return std::make_pair(std::move(arg_shardings), std::move(retval_shardings)); in ComputeArgAndRetvalShardings()168 const std::map<int, xla::OpSharding>& retval_shardings, in BuildComputation() argument224 auto it = retval_shardings.find(i); in BuildComputation()226 it == retval_shardings.end() ? absl::optional<xla::OpSharding>() in BuildComputation()228 if (it != retval_shardings.end()) { in BuildComputation()239 if (it != retval_shardings.end()) { in BuildComputation()1328 std::map<int, xla::OpSharding> retval_shardings; in CompileGraph() local1329 TF_ASSIGN_OR_RETURN(std::tie(arg_shardings, retval_shardings), in CompileGraph()[all …]
436 const std::vector<::xla::OpSharding>& retval_shardings,
2893 const std::vector<xla::OpSharding>& retval_shardings, in BuildExecuteNodes() argument3018 for (int i = 0; i < retval_shardings.size(); ++i) { in BuildExecuteNodes()3019 const auto& sharding = retval_shardings[i]; in BuildExecuteNodes()3283 const auto& sharding = retval_shardings[core_retval_nums[core][i]]; in BuildExecuteNodes()