| 113 | } // namespace |
| 114 | |
| 115 | XlaOp DynamicSliceInMinorDims(XlaOp x, absl::Span<const XlaOp> starts, |
| 116 | absl::Span<const int64> sizes) { |
| 117 | XlaBuilder* builder = x.builder(); |
| 118 | return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 119 | TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); |
| 120 | const int64 n_dims = shape.rank(); |
| 121 | int64 n_minor_dims = starts.size(); |
| 122 | TF_RET_CHECK(n_minor_dims == sizes.size()); |
| 123 | TF_RET_CHECK(n_minor_dims <= n_dims); |
| 124 | auto major_dims = AsInt64Slice(shape.dimensions()) |
| 125 | .subspan( |
| 126 | /*pos=*/0, |
| 127 | /*len=*/n_dims - sizes.size()); |
| 128 | TF_ASSIGN_OR_RETURN(auto padded_starts, PrependZerosInMajorDims(x, starts)); |
| 129 | auto padded_sizes = ConcatVectors(major_dims, sizes); |
| 130 | return DynamicSlice(x, padded_starts, padded_sizes); |
| 131 | }); |
| 132 | } |
| 133 | |
| 134 | XlaOp DynamicUpdateSliceInMinorDims(XlaOp x, XlaOp update, |
| 135 | absl::Span<const XlaOp> starts) { |