| 1033 | ) |
| 1034 | |
| 1035 | def __dask_annotations__(self) -> dict[str, dict[Key, object]]: |
| 1036 | # optimization has to be called (and cached) since blockwise fusion can |
| 1037 | # alter annotations |
| 1038 | # see `dask.blockwise.(_fuse_annotations|_can_fuse_annotations)` |
| 1039 | dsk = self._optimized_dsk |
| 1040 | annotations_by_type: defaultdict[str, dict[Key, object]] = defaultdict(dict) |
| 1041 | for layer in dsk.layers.values(): |
| 1042 | if layer.annotations: |
| 1043 | annot = layer.annotations |
| 1044 | for annot_type, value in annot.items(): |
| 1045 | annotations_by_type[annot_type].update( |
| 1046 | {k: (value(k) if callable(value) else value) for k in layer} |
| 1047 | ) |
| 1048 | return dict(annotations_by_type) |
| 1049 | |
| 1050 | def __dask_keys__(self): |
| 1051 | if (keys := self.operand("output_keys")) is not None: |