| 725 | ) |
| 726 | |
| 727 | def cull( |
| 728 | self, keys: set, all_hlg_keys: Iterable |
| 729 | ) -> tuple[Layer, Mapping[Key, set[Key]]]: |
| 730 | # Culling is simple for Blockwise layers. We can just |
| 731 | # collect a set of required output blocks (tuples), and |
| 732 | # only construct graph for these blocks in `make_blockwise_graph` |
| 733 | |
| 734 | output_blocks: set[tuple[int, ...]] = set() |
| 735 | for key in keys: |
| 736 | if key[0] == self.output: |
| 737 | output_blocks.add(key[1:]) |
| 738 | culled_deps = self._cull_dependencies(output_blocks) |
| 739 | out_size_iter = (self.dims[i] for i in self.output_indices) |
| 740 | if prod(out_size_iter) != len(culled_deps): |
| 741 | culled_layer = self._cull(output_blocks) |
| 742 | return culled_layer, culled_deps |
| 743 | else: |
| 744 | return self, culled_deps |
| 745 | |
| 746 | def clone( |
| 747 | self, |