(idtype, dtype)
| 218 | @parametrize_idtype |
| 219 | @pytest.mark.parametrize("dtype", [F.float32, F.float64]) |
| 220 | def test_csrmask_backward(idtype, dtype): |
| 221 | a, A = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, "A", "B", "AB") |
| 222 | b, B = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, "A", "B", "AB") |
| 223 | A_row, A_col = A.edges(order="eid") |
| 224 | B_row, B_col = B.edges(order="eid") |
| 225 | A_row = F.asnumpy(A_row) |
| 226 | A_col = F.asnumpy(A_col) |
| 227 | B_row = F.asnumpy(B_row) |
| 228 | B_col = F.asnumpy(B_col) |
| 229 | a_dense = F.attach_grad(F.tensor(a.todense(), dtype=dtype)) |
| 230 | |
| 231 | A.edata["w"] = F.attach_grad(A.edata["w"]) |
| 232 | |
| 233 | with F.record_grad(): |
| 234 | # Test for two element case |
| 235 | C1 = F.csrmask(A._graph, A.edata["w"], B._graph) |
| 236 | if dgl.backend.backend_name == "tensorflow": |
| 237 | import tensorflow as tf |
| 238 | |
| 239 | C2 = tf.gather_nd(a_dense, tf.stack([B_row, B_col], 1)) |
| 240 | else: |
| 241 | C2 = a_dense[B_row, B_col] |
| 242 | assert F.allclose(C1, C2, rtol=1e-4, atol=1e-4) |
| 243 | |
| 244 | F.backward(F.reduce_sum(C1) + F.reduce_sum(C2)) |
| 245 | a_dense_grad = F.asnumpy(F.grad(a_dense))[A_row, A_col] |
| 246 | A_spspmm_grad = F.asnumpy(F.grad(A.edata["w"])) |
| 247 | assert np.allclose(a_dense_grad, A_spspmm_grad, rtol=1e-4, atol=1e-4) |
| 248 | |
| 249 | |
| 250 | if __name__ == "__main__": |
no test coverage detected