| 766 | } |
| 767 | |
| 768 | std::vector<array> vmap_replace( |
| 769 | const std::vector<array>& inputs, |
| 770 | const std::vector<array>& s_inputs, |
| 771 | const std::vector<array>& s_outputs, |
| 772 | const std::vector<int>& in_axes, |
| 773 | const std::vector<int>& out_axes) { |
| 774 | if (out_axes.size() != s_outputs.size()) { |
| 775 | std::stringstream msg; |
| 776 | msg << "[vmap] The number of out axes (" << out_axes.size() |
| 777 | << ") must match the number of outputs (" << s_outputs.size() << ")."; |
| 778 | throw std::invalid_argument(msg.str()); |
| 779 | } |
| 780 | |
| 781 | int vmap_size = -1; |
| 782 | for (int i = 0; i < inputs.size(); ++i) { |
| 783 | if (in_axes[i] >= 0) { |
| 784 | vmap_size = inputs[i].shape(in_axes[i]); |
| 785 | break; |
| 786 | } |
| 787 | } |
| 788 | if (vmap_size == -1) { |
| 789 | throw std::invalid_argument("At least one of in_axes must be non-None."); |
| 790 | } |
| 791 | |
| 792 | std::unordered_map<std::uintptr_t, std::pair<array, int>> tmap; |
| 793 | std::unordered_set<std::uintptr_t> needs_vmap; |
| 794 | std::unordered_set<std::uintptr_t> cache; |
| 795 | for (int i = 0; i < s_inputs.size(); ++i) { |
| 796 | auto in = s_inputs[i]; |
| 797 | if (in_axes[i] != -1) { |
| 798 | tmap.insert({in.id(), {inputs[i], in_axes[i]}}); |
| 799 | needs_vmap.insert(in.id()); |
| 800 | in.set_tracer(false); |
| 801 | } |
| 802 | cache.insert(in.id()); |
| 803 | } |
| 804 | |
| 805 | // Topologically sort the graph |
| 806 | std::vector<array> tape; |
| 807 | |
| 808 | std::function<void(const array&)> recurse; |
| 809 | |
| 810 | recurse = [&](const array& a) { |
| 811 | auto id = a.id(); |
| 812 | if (cache.find(id) != cache.end()) { |
| 813 | return; |
| 814 | } |
| 815 | cache.insert(id); |
| 816 | for (auto& s : a.siblings()) { |
| 817 | cache.insert(s.id()); |
| 818 | } |
| 819 | |
| 820 | // Recurse on inputs |
| 821 | for (auto& input : a.inputs()) { |
| 822 | recurse(input); |
| 823 | } |
| 824 | // If any input needs a vmap, then the outputs also need |
| 825 | // a vmap |