| 898 | } // namespace detail |
| 899 | |
| 900 | std::function<std::vector<array>(const std::vector<array>&)> vmap( |
| 901 | const std::function<std::vector<array>(const std::vector<array>&)>& fun, |
| 902 | const std::vector<int>& in_axes /* = {} */, |
| 903 | const std::vector<int>& out_axes /* = {} */) { |
| 904 | auto infer_axes = [](auto axes) { |
| 905 | return !axes.empty() && |
| 906 | std::all_of(axes.begin(), axes.end(), [](int ax) { return ax < 0; }); |
| 907 | }; |
| 908 | if (infer_axes(in_axes) != infer_axes(out_axes)) { |
| 909 | throw std::invalid_argument( |
| 910 | "[vmap] Input (or output) axes must be " |
| 911 | "specified if output (or input) axes are."); |
| 912 | } |
| 913 | auto vfun = [fun, in_axes = in_axes, out_axes = out_axes]( |
| 914 | const std::vector<array>& inputs) mutable { |
| 915 | if (in_axes.size() == 0) { |
| 916 | in_axes.resize(inputs.size(), 0); |
| 917 | } |
| 918 | |
| 919 | auto [trace_inputs, trace_outputs] = |
| 920 | detail::vmap_trace(fun, inputs, in_axes); |
| 921 | |
| 922 | if (out_axes.size() == 0) { |
| 923 | out_axes.resize(trace_outputs.size(), 0); |
| 924 | } |
| 925 | |
| 926 | return detail::vmap_replace( |
| 927 | inputs, trace_inputs, trace_outputs, in_axes, out_axes); |
| 928 | }; |
| 929 | |
| 930 | return vfun; |
| 931 | } |
| 932 | |
| 933 | std::function<array(const array&, const array&)> vmap( |
| 934 | const std::function<array(const array&, const array&)>& fun, |
no test coverage detected