| 22 | namespace { |
| 23 | |
| 24 | std::tuple<array, array, int> vmap_binary_op( |
| 25 | const std::vector<array>& inputs, |
| 26 | const std::vector<int>& axes, |
| 27 | const Stream& stream) { |
| 28 | assert(inputs.size() == 2); |
| 29 | assert(axes.size() == 2); |
| 30 | |
| 31 | if (axes[0] == -1 && axes[1] == -1) { |
| 32 | return {inputs[0], inputs[1], -1}; |
| 33 | } |
| 34 | |
| 35 | auto a = inputs[0]; |
| 36 | auto b = inputs[1]; |
| 37 | int ndim = std::max(a.ndim() + (axes[0] == -1), b.ndim() + (axes[1] == -1)); |
| 38 | |
| 39 | auto expand_dims = [stream, ndim](auto in) { |
| 40 | auto shape = in.shape(); |
| 41 | shape.insert(shape.begin(), ndim - shape.size(), 1); |
| 42 | return reshape(in, shape, stream); |
| 43 | }; |
| 44 | |
| 45 | int to_ax = (ndim - a.ndim()) + axes[0]; |
| 46 | int from_ax = (ndim - b.ndim()) + axes[1]; |
| 47 | a = expand_dims(a); |
| 48 | b = expand_dims(b); |
| 49 | |
| 50 | if (from_ax != to_ax) { |
| 51 | std::vector<int> tdims(b.ndim()); |
| 52 | std::iota(tdims.begin(), tdims.end(), 0); |
| 53 | tdims.erase(tdims.begin() + from_ax); |
| 54 | tdims.insert(tdims.begin() + to_ax, from_ax); |
| 55 | b = transpose(b, tdims, stream); |
| 56 | } |
| 57 | return {a, b, to_ax}; |
| 58 | } |
| 59 | |
| 60 | std::tuple<array, array, array, int> vmap_ternary_op( |
| 61 | const std::vector<array>& inputs, |