| 603 | } |
| 604 | |
| 605 | std::pair<std::vector<int64>, std::vector<int64>> OpTest::BroadcastableDims( |
| 606 | std::vector<int64> dims) { |
| 607 | if (dims.empty()) return {dims, dims}; |
| 608 | |
| 609 | // Remove some dimensions from the front of 'dims'. |
| 610 | size_t skip = |
| 611 | std::uniform_int_distribution<size_t>(0, dims.size() - 1)(generator()); |
| 612 | |
| 613 | std::vector<int64> bdims(dims.begin() + skip, dims.end()); |
| 614 | |
| 615 | // Randomly replace some of the remaining dimensions of 'dims' with 1. |
| 616 | std::bernoulli_distribution random_bool; |
| 617 | |
| 618 | for (int64& dim : bdims) { |
| 619 | if (random_bool(generator())) { |
| 620 | dim = 1LL; |
| 621 | } |
| 622 | } |
| 623 | |
| 624 | // Possibly swap the roles of 'dims' and 'bdims'. |
| 625 | if (random_bool(generator())) { |
| 626 | dims.swap(bdims); |
| 627 | } |
| 628 | return {dims, bdims}; |
| 629 | } |
| 630 | |
| 631 | std::pair<std::vector<int64>, std::vector<int64>> OpTest::BroadcastableDims() { |
| 632 | return BroadcastableDims(RandomDims(0, 3)); |