| 204 | def _softmax_api_helper(self, dev): |
| 205 | |
| 206 | def _run_test(dev, org_shape, axis, aft_shape): |
| 207 | x_0 = np.random.random(org_shape).astype(np.float32) |
| 208 | x_0 = x_0 + 1000 |
| 209 | x0 = tensor.Tensor(device=dev, data=x_0) |
| 210 | |
| 211 | # test with axis |
| 212 | y0 = tensor._call_singa_func(singa_api.SoftMax, x0.data, axis) |
| 213 | |
| 214 | # test with numpy |
| 215 | x_0 = x_0.reshape(aft_shape) |
| 216 | x_0 = x_0 - np.max(x_0) |
| 217 | y1 = np.divide(np.exp(x_0), |
| 218 | np.sum(np.exp(x_0), axis=1).reshape(x_0.shape[0], |
| 219 | 1)) # 2d softmax |
| 220 | y1 = y1.reshape(org_shape) |
| 221 | |
| 222 | np.testing.assert_array_almost_equal(tensor.to_numpy(y0), y1) |
| 223 | |
| 224 | _run_test(dev, [2, 2], 1, [2, 2]) |
| 225 | _run_test(dev, [2, 2], 0, [1, 4]) |