| 80 | } |
| 81 | |
| 82 | StatusOr<std::vector<XlaOp>> ForEachIndex( |
| 83 | int64 num_iterations, PrimitiveType num_iterations_type, |
| 84 | const ForEachIndexBodyFunction& body_function, |
| 85 | absl::Span<const XlaOp> initial_values, absl::string_view name, |
| 86 | XlaBuilder* builder) { |
| 87 | auto while_cond_fn = [&](absl::Span<const XlaOp> values, |
| 88 | XlaBuilder* cond_builder) -> StatusOr<XlaOp> { |
| 89 | return Lt(values[0], ConstantR0WithType(cond_builder, num_iterations_type, |
| 90 | num_iterations)); |
| 91 | }; |
| 92 | auto while_body_fn = |
| 93 | [&](absl::Span<const XlaOp> values, |
| 94 | XlaBuilder* body_builder) -> StatusOr<std::vector<XlaOp>> { |
| 95 | XlaOp iteration = values[0]; |
| 96 | |
| 97 | std::vector<XlaOp> updated_values; |
| 98 | updated_values.reserve(values.size()); |
| 99 | updated_values.push_back(Add( |
| 100 | iteration, |
| 101 | ConstantLiteral(body_builder, LiteralUtil::One(num_iterations_type)))); |
| 102 | |
| 103 | values.remove_prefix(1); |
| 104 | TF_ASSIGN_OR_RETURN(std::vector<XlaOp> body_outputs, |
| 105 | body_function(iteration, values, body_builder)); |
| 106 | updated_values.insert(updated_values.end(), body_outputs.begin(), |
| 107 | body_outputs.end()); |
| 108 | return updated_values; |
| 109 | }; |
| 110 | |
| 111 | std::vector<XlaOp> values; |
| 112 | values.reserve(initial_values.size() + 1); |
| 113 | values.push_back( |
| 114 | ConstantLiteral(builder, LiteralUtil::Zero(num_iterations_type))); |
| 115 | values.insert(values.end(), initial_values.begin(), initial_values.end()); |
| 116 | |
| 117 | TF_ASSIGN_OR_RETURN(values, WhileLoopHelper(while_cond_fn, while_body_fn, |
| 118 | values, name, builder)); |
| 119 | values.erase(values.begin(), values.begin() + 1); |
| 120 | return values; |
| 121 | } |
| 122 | |
| 123 | } // namespace xla |