(sparse)
| 479 | |
| 480 | |
| 481 | def reduce_sparse_tensor(sparse): |
| 482 | if sparse.layout is torch.sparse_coo: |
| 483 | rebuild_indices_func, rebuild_indices_args = reduce_tensor(sparse._indices()) |
| 484 | rebuild_values_func, rebuild_values_args = reduce_tensor(sparse._values()) |
| 485 | return ( |
| 486 | rebuild_sparse_coo_tensor, |
| 487 | ( |
| 488 | rebuild_indices_func, |
| 489 | rebuild_indices_args, |
| 490 | rebuild_values_func, |
| 491 | rebuild_values_args, |
| 492 | sparse.shape, |
| 493 | sparse.is_coalesced(), |
| 494 | ), |
| 495 | ) |
| 496 | else: |
| 497 | if sparse.layout in {torch.sparse_csr, torch.sparse_bsr}: |
| 498 | compressed_indices = sparse.crow_indices() |
| 499 | plain_indices = sparse.col_indices() |
| 500 | elif sparse.layout in {torch.sparse_csc, torch.sparse_bsc}: |
| 501 | compressed_indices = sparse.ccol_indices() |
| 502 | plain_indices = sparse.row_indices() |
| 503 | else: |
| 504 | raise NotImplementedError(sparse.layout) |
| 505 | ( |
| 506 | rebuild_compressed_indices_func, |
| 507 | rebuild_compressed_indices_args, |
| 508 | ) = reduce_tensor(compressed_indices) |
| 509 | rebuild_plain_indices_func, rebuild_plain_indices_args = reduce_tensor( |
| 510 | plain_indices |
| 511 | ) |
| 512 | rebuild_values_func, rebuild_values_args = reduce_tensor(sparse.values()) |
| 513 | return ( |
| 514 | rebuild_sparse_compressed_tensor, |
| 515 | ( |
| 516 | rebuild_compressed_indices_func, |
| 517 | rebuild_compressed_indices_args, |
| 518 | rebuild_plain_indices_func, |
| 519 | rebuild_plain_indices_args, |
| 520 | rebuild_values_func, |
| 521 | rebuild_values_args, |
| 522 | sparse.shape, |
| 523 | sparse.layout, |
| 524 | ), |
| 525 | ) |
| 526 | |
| 527 | |
| 528 | def fd_id(fd): |
no test coverage detected