()
| 1799 | |
| 1800 | |
| 1801 | def test_batchnorm2d(): |
| 1802 | class BatchNorm2d1(Module): |
| 1803 | def __init__(self): |
| 1804 | super().__init__() |
| 1805 | self.bn = torch.nn.BatchNorm2d(3) |
| 1806 | |
| 1807 | def forward(self, input): |
| 1808 | return self.bn(input) |
| 1809 | |
| 1810 | @tvm.script.ir_module |
| 1811 | class expected1: |
| 1812 | @R.function |
| 1813 | def main( |
| 1814 | input_1: R.Tensor((1, 3, 10, 10), dtype="float32"), |
| 1815 | w1: R.Tensor((3,), dtype="float32"), |
| 1816 | w2: R.Tensor((3,), dtype="float32"), |
| 1817 | w3: R.Tensor((3,), dtype="float32"), |
| 1818 | w4: R.Tensor((3,), dtype="float32"), |
| 1819 | ) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")): |
| 1820 | # block 0 |
| 1821 | with R.dataflow(): |
| 1822 | lv: R.Tuple( |
| 1823 | R.Tensor((1, 3, 10, 10), dtype="float32"), |
| 1824 | R.Tensor((3,), dtype="float32"), |
| 1825 | R.Tensor((3,), dtype="float32"), |
| 1826 | ) = R.nn.batch_norm( |
| 1827 | input_1, |
| 1828 | w1, |
| 1829 | w2, |
| 1830 | w3, |
| 1831 | w4, |
| 1832 | axis=1, |
| 1833 | epsilon=1e-05, |
| 1834 | center=True, |
| 1835 | scale=True, |
| 1836 | momentum=0.1, |
| 1837 | training=False, |
| 1838 | ) |
| 1839 | lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = lv[0] |
| 1840 | gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv1,) |
| 1841 | R.output(gv) |
| 1842 | return gv |
| 1843 | |
| 1844 | class BatchNorm2dCustom(Module): |
| 1845 | def __init__(self): |
| 1846 | super().__init__() |
| 1847 | self.bn = torch.nn.BatchNorm2d(3, eps=0.001, momentum=0.01) |
| 1848 | |
| 1849 | def forward(self, input): |
| 1850 | return self.bn(input) |
| 1851 | |
| 1852 | @tvm.script.ir_module |
| 1853 | class expected2: |
| 1854 | @R.function |
| 1855 | def main( |
| 1856 | input_1: R.Tensor((1, 3, 10, 10), dtype="float32"), |
| 1857 | w1: R.Tensor((3,), dtype="float32"), |
| 1858 | w2: R.Tensor((3,), dtype="float32"), |
nothing calls this directly
no test coverage detected
searching dependent graphs…