| 48 | } |
| 49 | |
| 50 | xla::StatusOr<xla::XlaOp> Expand(xla::XlaOp input, int64 dim) { |
| 51 | xla::XlaBuilder* builder = input.builder(); |
| 52 | TF_ASSIGN_OR_RETURN(xla::Shape input_shape, builder->GetShape(input)); |
| 53 | |
| 54 | if (input_shape.dimensions(dim) % 4 != 0) { |
| 55 | return errors::InvalidArgument( |
| 56 | "Expected vectorized dimension to be evenly divisible by 4; got ", |
| 57 | input_shape.dimensions(dim)); |
| 58 | } |
| 59 | |
| 60 | // Split the `dim` into two dimensions with a reshape. The size of the new |
| 61 | // dimension is always 4. |
| 62 | std::vector<int64> expanded_shape = |
| 63 | xla::SpanToVector(input_shape.dimensions()); |
| 64 | expanded_shape[dim] /= 4; |
| 65 | expanded_shape.insert(expanded_shape.begin() + dim, 4); |
| 66 | |
| 67 | // Move the newly created dimension to the end with a transpose. |
| 68 | std::vector<int64> permutation; |
| 69 | for (int64 i = 0; i != expanded_shape.size(); ++i) { |
| 70 | permutation.push_back(i); |
| 71 | if (i == dim) { |
| 72 | ++i; |
| 73 | } |
| 74 | } |
| 75 | permutation.push_back(dim + 1); |
| 76 | |
| 77 | return xla::Transpose(xla::Reshape(input, expanded_shape), permutation); |
| 78 | } |
| 79 | |
| 80 | } // namespace |
| 81 |
no test coverage detected