(self, dev)
| 1829 | # self._split_helper(gpu_dev) |
| 1830 | |
| 1831 | def _gather_helper(self, dev): |
| 1832 | X = np.array([0, 1, 2]).astype(np.float32) |
| 1833 | x = tensor.from_numpy(X) |
| 1834 | x.to_device(dev) |
| 1835 | y = autograd.gather(x, 0, [0, 1, 3]) |
| 1836 | |
| 1837 | # frontend |
| 1838 | model = sonnx.to_onnx([x], [y]) |
| 1839 | # print('The model is:\n{}'.format(model)) |
| 1840 | |
| 1841 | # backend |
| 1842 | sg_ir = sonnx.prepare(model, device=dev) |
| 1843 | sg_ir.is_graph = True |
| 1844 | y_t = sg_ir.run([x]) |
| 1845 | |
| 1846 | np.testing.assert_array_almost_equal( |
| 1847 | tensor.to_numpy(y).shape, |
| 1848 | tensor.to_numpy(y_t[0]).shape) |
| 1849 | |
| 1850 | def test_gather_cpu(self): |
| 1851 | self._gather_helper(cpu_dev) |
no test coverage detected