| 3808 | } |
| 3809 | |
| 3810 | array softmax( |
| 3811 | const array& a, |
| 3812 | const std::vector<int>& axes, |
| 3813 | bool precise /* = false */, |
| 3814 | StreamOrDevice s /* = {}*/) { |
| 3815 | if (a.size() == 0) { |
| 3816 | return a; |
| 3817 | } |
| 3818 | if (a.ndim() == 0 && !axes.empty()) { |
| 3819 | throw std::invalid_argument( |
| 3820 | "[softmax] Received non-empty axes for array with 0 dimensions."); |
| 3821 | } |
| 3822 | bool reduce_last_dim = |
| 3823 | !axes.empty() && (axes.back() == a.ndim() - 1 || axes.back() == -1); |
| 3824 | if (reduce_last_dim) { |
| 3825 | // For more than 2 axes check if axes is [0, 1, ..., NDIM - 1] and shape |
| 3826 | // is [1, 1, ..., N]. |
| 3827 | for (int i = axes.size() - 2; i >= 0; --i) { |
| 3828 | if ((axes[i] + 1 != axes[i + 1]) || (a.shape(axes[i]) != 1)) { |
| 3829 | reduce_last_dim = false; |
| 3830 | break; |
| 3831 | } |
| 3832 | } |
| 3833 | } |
| 3834 | bool is_complex = issubdtype(a.dtype(), complexfloating); |
| 3835 | if (!is_complex && reduce_last_dim) { |
| 3836 | auto dtype = at_least_float(a.dtype()); |
| 3837 | return array( |
| 3838 | a.shape(), |
| 3839 | dtype, |
| 3840 | std::make_shared<Softmax>(to_stream(s), precise), |
| 3841 | {astype(a, dtype, s)}); |
| 3842 | } else { |
| 3843 | auto in = a; |
| 3844 | if (precise && !is_complex) { |
| 3845 | in = astype(a, float32, s); |
| 3846 | } |
| 3847 | auto a_max = stop_gradient(max(in, axes, /*keepdims = */ true, s), s); |
| 3848 | auto ex = exp(subtract(in, a_max, s), s); |
| 3849 | return astype( |
| 3850 | divide(ex, sum(ex, axes, /*keepdims = */ true, s), s), a.dtype(), s); |
| 3851 | } |
| 3852 | } |
| 3853 | |
| 3854 | array softmax( |
| 3855 | const array& a, |
nothing calls this directly
no test coverage detected