| 958 | } |
| 959 | |
| 960 | std::function<std::vector<array>(const std::vector<array>&)> custom_function( |
| 961 | std::function<std::vector<array>(const std::vector<array>&)> fun, |
| 962 | std::optional<std::function<std::vector<array>( |
| 963 | const std::vector<array>&, |
| 964 | const std::vector<array>&, |
| 965 | const std::vector<array>&)>> fun_vjp /* = std::nullopt */, |
| 966 | std::optional<std::function<std::vector<array>( |
| 967 | const std::vector<array>&, |
| 968 | const std::vector<array>&, |
| 969 | const std::vector<int>&)>> fun_jvp /* = std::nullopt */, |
| 970 | std::optional<std::function<std::pair<std::vector<array>, std::vector<int>>( |
| 971 | const std::vector<array>&, |
| 972 | const std::vector<int>&)>> fun_vmap /* = std::nullopt */) { |
| 973 | if (!fun_vjp.has_value() && !fun_jvp.has_value() && !fun_vmap.has_value()) { |
| 974 | return fun; |
| 975 | } |
| 976 | |
| 977 | return [fun = std::move(fun), |
| 978 | fun_vjp = std::move(fun_vjp), |
| 979 | fun_jvp = std::move(fun_jvp), |
| 980 | fun_vmap = std::move(fun_vmap)](const std::vector<array>& args) { |
| 981 | // Compute the outputs |
| 982 | auto outputs = fun(args); |
| 983 | for (auto& out : outputs) { |
| 984 | out = stop_gradient(out); |
| 985 | } |
| 986 | |
| 987 | // Prepare the inputs to the primitive |
| 988 | // We also add the outputs to the primitive so that it can "run" the forward |
| 989 | // pass. |
| 990 | std::vector<array> inputs = args; |
| 991 | inputs.insert(inputs.end(), outputs.begin(), outputs.end()); |
| 992 | |
| 993 | // Compute the stream. Maybe do it in a smarter way at some point in the |
| 994 | // future. |
| 995 | Stream s = (outputs[0].has_primitive()) ? outputs[0].primitive().stream() |
| 996 | : default_stream(default_device()); |
| 997 | |
| 998 | // Make the output info |
| 999 | std::vector<Shape> shapes; |
| 1000 | std::vector<Dtype> dtypes; |
| 1001 | for (const auto& out : outputs) { |
| 1002 | shapes.emplace_back(out.shape()); |
| 1003 | dtypes.emplace_back(out.dtype()); |
| 1004 | } |
| 1005 | |
| 1006 | return array::make_arrays( |
| 1007 | std::move(shapes), |
| 1008 | dtypes, |
| 1009 | std::make_shared<CustomTransforms>( |
| 1010 | to_stream(s), |
| 1011 | outputs.size(), |
| 1012 | |
| 1013 | // We use the passed vjp function or compute it from the inputs and |
| 1014 | // passed cotangents. Note that this may be less efficient than |
| 1015 | // using `fun` directly because we may not be able to fully reuse |
| 1016 | // the outputs of the forward pass. |
| 1017 | fun_vjp.value_or( |
no test coverage detected