(self, dev)
| 3012 | self.split_test(gpu_dev) |
| 3013 | |
| 3014 | def gather_test(self, dev): |
| 3015 | config = [([0, 1, 3], 0), ([0, 1, 3], 1), ([[0, 1], [1, 2], [2, 3]], 1), |
| 3016 | ([0, -1, -2], 0)] # (indices, axis) |
| 3017 | for indices, _axis in config: |
| 3018 | X = np.random.randn(5, 4, 3, 2).astype(np.float32) |
| 3019 | y = np.take(X, indices, axis=_axis) |
| 3020 | DY = np.ones(y.shape, dtype=np.float32) |
| 3021 | |
| 3022 | x = tensor.from_numpy(X) |
| 3023 | dy = tensor.from_numpy(DY) |
| 3024 | x.to_device(dev) |
| 3025 | dy.to_device(dev) |
| 3026 | |
| 3027 | result = autograd.gather(x, _axis, indices) |
| 3028 | dx = result.creator.backward(dy.data) |
| 3029 | |
| 3030 | np.testing.assert_array_almost_equal(tensor.to_numpy(result), |
| 3031 | y, |
| 3032 | decimal=5) |
| 3033 | self.check_shape(dx.shape(), tuple(X.shape)) |
| 3034 | |
| 3035 | def test_gather_cpu(self): |
| 3036 | self.gather_test(cpu_dev) |
no test coverage detected