| 56 | } |
| 57 | |
| 58 | XlaOp UpdateSlice(XlaOp x, XlaOp update, absl::Span<const int64> start) { |
| 59 | XlaBuilder* builder = x.builder(); |
| 60 | return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 61 | TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); |
| 62 | const int64 n_dims = shape.rank(); |
| 63 | TF_RET_CHECK(start.size() == n_dims); |
| 64 | |
| 65 | // TODO(phawkins): make int64 work on all backends, remove the int32 cast. |
| 66 | std::vector<int32> start_as_int32(start.begin(), start.end()); |
| 67 | std::vector<XlaOp> start_ops(start.size()); |
| 68 | for (int i = 0; i < start.size(); ++i) { |
| 69 | start_ops[i] = ConstantR0(builder, start_as_int32[i]); |
| 70 | } |
| 71 | return DynamicUpdateSlice(x, update, start_ops); |
| 72 | }); |
| 73 | } |
| 74 | |
| 75 | XlaOp UpdateSliceInMinorDims(XlaOp x, XlaOp update, |
| 76 | absl::Span<const int64> start) { |
no test coverage detected