(idtype, dtype, nelems)
| 134 | @pytest.mark.parametrize("dtype", [F.float32, F.float64]) |
| 135 | @pytest.mark.parametrize("nelems", [1, 2]) |
| 136 | def test_csrsum_backward(idtype, dtype, nelems): |
| 137 | a, A = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, "A", "B", "AB") |
| 138 | b, B = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, "A", "B", "AB") |
| 139 | A_row, A_col = A.edges(order="eid") |
| 140 | B_row, B_col = B.edges(order="eid") |
| 141 | A_row = F.asnumpy(A_row) |
| 142 | A_col = F.asnumpy(A_col) |
| 143 | B_row = F.asnumpy(B_row) |
| 144 | B_col = F.asnumpy(B_col) |
| 145 | a_dense = F.attach_grad(F.tensor(a.todense(), dtype=dtype)) |
| 146 | b_dense = F.attach_grad(F.tensor(b.todense(), dtype=dtype)) |
| 147 | |
| 148 | A.edata["w"] = F.attach_grad(A.edata["w"]) |
| 149 | B.edata["w"] = F.attach_grad(B.edata["w"]) |
| 150 | |
| 151 | with F.record_grad(): |
| 152 | if nelems == 2: |
| 153 | # Test for two element case |
| 154 | C = dgl.adj_sum_graph([A, B], "w") |
| 155 | assert C.canonical_etypes == A.canonical_etypes |
| 156 | C_dense = np.zeros((3, 4)) |
| 157 | C_row, C_col = C.edges(order="eid") |
| 158 | C_row = F.asnumpy(C_row) |
| 159 | C_col = F.asnumpy(C_col) |
| 160 | C_dense[C_row, C_col] = F.asnumpy(C.edata["w"]) |
| 161 | c_dense = a_dense + b_dense |
| 162 | assert np.allclose( |
| 163 | C_dense, F.asnumpy(c_dense), rtol=1e-4, atol=1e-4 |
| 164 | ) |
| 165 | |
| 166 | F.backward(F.reduce_sum(C.edata["w"]) + F.reduce_sum(c_dense)) |
| 167 | a_dense_grad = F.asnumpy(F.grad(a_dense))[A_row, A_col] |
| 168 | b_dense_grad = F.asnumpy(F.grad(b_dense))[B_row, B_col] |
| 169 | A_spspmm_grad = F.asnumpy(F.grad(A.edata["w"])) |
| 170 | B_spspmm_grad = F.asnumpy(F.grad(B.edata["w"])) |
| 171 | assert np.allclose( |
| 172 | a_dense_grad, A_spspmm_grad, rtol=1e-4, atol=1e-4 |
| 173 | ) |
| 174 | assert np.allclose( |
| 175 | b_dense_grad, B_spspmm_grad, rtol=1e-4, atol=1e-4 |
| 176 | ) |
| 177 | elif nelems == 1: |
| 178 | # Test for single element case |
| 179 | C = dgl.adj_sum_graph([A], "w") |
| 180 | assert C.canonical_etypes == A.canonical_etypes |
| 181 | C_dense = np.zeros((3, 4)) |
| 182 | C_row, C_col = C.edges(order="eid") |
| 183 | C_row = F.asnumpy(C_row) |
| 184 | C_col = F.asnumpy(C_col) |
| 185 | C_dense[C_row, C_col] = F.asnumpy(C.edata["w"]) |
| 186 | c_dense = a_dense |
| 187 | assert np.allclose( |
| 188 | C_dense, F.asnumpy(c_dense), rtol=1e-4, atol=1e-4 |
| 189 | ) |
| 190 | |
| 191 | F.backward(F.reduce_sum(C.edata["w"]) + F.reduce_sum(c_dense)) |
| 192 | a_dense_grad = F.asnumpy(F.grad(a_dense))[A_row, A_col] |
| 193 | A_spspmm_grad = F.asnumpy(F.grad(A.edata["w"])) |
no test coverage detected