(idtype)
| 1802 | |
| 1803 | @parametrize_idtype |
| 1804 | def test_apply(idtype): |
| 1805 | def node_udf(nodes): |
| 1806 | return {"h": nodes.data["h"] * 2} |
| 1807 | |
| 1808 | def node_udf2(nodes): |
| 1809 | return {"h": F.sum(nodes.data["h"], dim=1, keepdims=True)} |
| 1810 | |
| 1811 | def edge_udf(edges): |
| 1812 | return {"h": edges.data["h"] * 2 + edges.src["h"]} |
| 1813 | |
| 1814 | g = create_test_heterograph(idtype) |
| 1815 | g.nodes["user"].data["h"] = F.ones((3, 5)) |
| 1816 | g.apply_nodes(node_udf, ntype="user") |
| 1817 | assert F.array_equal(g.nodes["user"].data["h"], F.ones((3, 5)) * 2) |
| 1818 | |
| 1819 | g["plays"].edata["h"] = F.ones((4, 5)) |
| 1820 | g.apply_edges(edge_udf, etype=("user", "plays", "game")) |
| 1821 | assert F.array_equal(g["plays"].edata["h"], F.ones((4, 5)) * 4) |
| 1822 | |
| 1823 | # test apply on graph with only one type |
| 1824 | g["follows"].apply_nodes(node_udf) |
| 1825 | assert F.array_equal(g.nodes["user"].data["h"], F.ones((3, 5)) * 4) |
| 1826 | |
| 1827 | g["plays"].apply_edges(edge_udf) |
| 1828 | assert F.array_equal(g["plays"].edata["h"], F.ones((4, 5)) * 12) |
| 1829 | |
| 1830 | # Test the case that feature size changes |
| 1831 | g.nodes["user"].data["h"] = F.ones((3, 5)) |
| 1832 | g.apply_nodes(node_udf2, ntype="user") |
| 1833 | assert F.array_equal(g.nodes["user"].data["h"], F.ones((3, 1)) * 5) |
| 1834 | |
| 1835 | # test fail case |
| 1836 | # fail due to multiple types |
| 1837 | with pytest.raises(DGLError): |
| 1838 | g.apply_nodes(node_udf) |
| 1839 | |
| 1840 | with pytest.raises(DGLError): |
| 1841 | g.apply_edges(edge_udf) |
| 1842 | |
| 1843 | |
| 1844 | @parametrize_idtype |
nothing calls this directly
no test coverage detected