| 73 | } |
| 74 | |
| 75 | XlaOp UpdateSliceInMinorDims(XlaOp x, XlaOp update, |
| 76 | absl::Span<const int64> start) { |
| 77 | XlaBuilder* builder = x.builder(); |
| 78 | return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 79 | TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); |
| 80 | const int64 n_dims = shape.rank(); |
| 81 | const int64 n_minor_dims = start.size(); |
| 82 | TF_RET_CHECK(n_minor_dims <= n_dims); |
| 83 | std::vector<int64> padded_start(n_dims, 0); |
| 84 | std::copy(start.begin(), start.end(), |
| 85 | padded_start.begin() + (n_dims - n_minor_dims)); |
| 86 | return UpdateSlice(x, update, padded_start); |
| 87 | }); |
| 88 | } |
| 89 | |
| 90 | namespace { |
| 91 |
no test coverage detected