| 257 | } |
| 258 | |
| 259 | Status XlaBuilder::SetDynamicBinding(int64 dynamic_size_param_num, |
| 260 | ShapeIndex dynamic_size_param_index, |
| 261 | int64 target_param_num, |
| 262 | ShapeIndex target_param_index, |
| 263 | int64 target_dim_num) { |
| 264 | bool param_exists = false; |
| 265 | for (size_t index = 0; index < instructions_.size(); ++index) { |
| 266 | HloInstructionProto& instr = instructions_[index]; |
| 267 | if (instr.opcode() == HloOpcodeString(HloOpcode::kParameter) && |
| 268 | instr.parameter_number() == target_param_num) { |
| 269 | param_exists = true; |
| 270 | Shape param_shape(instr.shape()); |
| 271 | Shape* param_shape_ptr = ¶m_shape; |
| 272 | for (int64 index : target_param_index) { |
| 273 | param_shape_ptr = param_shape_ptr->mutable_tuple_shapes(index); |
| 274 | } |
| 275 | param_shape_ptr->set_dynamic_dimension(target_dim_num, |
| 276 | /*is_dynamic=*/true); |
| 277 | *instr.mutable_shape() = param_shape.ToProto(); |
| 278 | instruction_shapes_[index] = |
| 279 | absl::make_unique<Shape>(std::move(param_shape)); |
| 280 | } |
| 281 | } |
| 282 | if (!param_exists) { |
| 283 | return InvalidArgument( |
| 284 | "Asked to mark parameter %lld as dynamic sized parameter, but the " |
| 285 | "doesn't exists", |
| 286 | target_param_num); |
| 287 | } |
| 288 | |
| 289 | TF_RETURN_IF_ERROR(dynamic_parameter_binding_.Bind( |
| 290 | DynamicParameterBinding::DynamicParameter{dynamic_size_param_num, |
| 291 | dynamic_size_param_index}, |
| 292 | DynamicParameterBinding::DynamicDimension{ |
| 293 | target_param_num, target_param_index, target_dim_num})); |
| 294 | return Status::OK(); |
| 295 | } |
| 296 | |
| 297 | Status XlaBuilder::SetInstructionFrontendAttribute(const XlaOp op, |
| 298 | std::string attribute, |