| 978 | return ret; |
| 979 | } |
| 980 | void SoftMaxBackward(const Tensor &in, Tensor *out, int axis, |
| 981 | const Tensor &fdout) { |
| 982 | // {a_0, a_1, ..., a_k-1, a_k, ... a_n-1} |
| 983 | // reshape to |
| 984 | // { a_0 * a_1 * ... a_k-1, a_k * ... a_n-1 } |
| 985 | |
| 986 | // assert axis \in {-r, r-1} |
| 987 | CHECK_LE(axis, (int)in.shape().size() - 1); |
| 988 | CHECK_GE(axis, -1 * (int)in.nDim()); |
| 989 | |
| 990 | Shape original_shape = in.shape(); |
| 991 | if (axis < 0) axis = in.shape().size() + axis; |
| 992 | |
| 993 | Shape coerced_shape = {1, 1}; |
| 994 | for (std::size_t i = 0, max = in.shape().size(); i != max; ++i) { |
| 995 | if (i < axis) |
| 996 | coerced_shape[0] *= in.shape()[i]; |
| 997 | else |
| 998 | coerced_shape[1] *= in.shape()[i]; |
| 999 | } |
| 1000 | |
| 1001 | Tensor in_reshaped = Reshape(in, coerced_shape); |
| 1002 | out->Reshape(coerced_shape); |
| 1003 | |
| 1004 | do { |
| 1005 | TYPE_LANG_SWITCH(in.data_type(), DType, in.device()->lang(), Lang, { |
| 1006 | Tensor &outRef = *out; |
| 1007 | out->device()->Exec( |
| 1008 | [in, outRef, fdout](Context *ctx) mutable { |
| 1009 | SoftMaxBackward<DType, Lang>(in, &outRef, fdout, ctx); |
| 1010 | }, |
| 1011 | {in.block(), fdout.block()}, {out->block()}, "SoftmaxBackward"); |
| 1012 | }); |
| 1013 | } while (0); |
| 1014 | |
| 1015 | out->Reshape(original_shape); |
| 1016 | } |
| 1017 | |
| 1018 | Tensor SoftMaxBackward(const Tensor &in, int axis, const Tensor &fdout) { |
| 1019 | Tensor ret(in.shape(), in.device(), in.data_type()); |