(nt)
| 421 | |
| 422 | |
| 423 | def reduce_nested_tensor(nt): |
| 424 | rebuild_buffer_func, rebuild_buffer_args = reduce_tensor(nt.values()) |
| 425 | rebuild_sizes_func, rebuild_sizes_args = reduce_tensor(nt._nested_tensor_size()) |
| 426 | rebuild_strides_func, rebuild_strides_args = reduce_tensor( |
| 427 | nt._nested_tensor_strides() |
| 428 | ) |
| 429 | rebuild_offsets_func, rebuild_offsets_args = reduce_tensor( |
| 430 | nt._nested_tensor_storage_offsets() |
| 431 | ) |
| 432 | |
| 433 | return ( |
| 434 | rebuild_nested_tensor, |
| 435 | ( |
| 436 | rebuild_buffer_func, |
| 437 | rebuild_buffer_args, |
| 438 | rebuild_sizes_func, |
| 439 | rebuild_sizes_args, |
| 440 | rebuild_strides_func, |
| 441 | rebuild_strides_args, |
| 442 | rebuild_offsets_func, |
| 443 | rebuild_offsets_args, |
| 444 | ), |
| 445 | ) |
| 446 | |
| 447 | |
| 448 | def rebuild_sparse_coo_tensor( |
no test coverage detected