Construct a chunks tuple that broadcasts many chunks tuples >>> a = ((5, 5),) >>> b = ((5, 5),) >>> broadcast_chunks(a, b) ((5, 5),) >>> a = ((10, 10, 10), (5, 5),) >>> b = ((5, 5),) >>> broadcast_chunks(a, b) ((10, 10, 10), (5, 5)) >>> a = ((10, 10, 10), (5, 5
(*chunkss)
| 1042 | |
| 1043 | |
| 1044 | def broadcast_chunks(*chunkss): |
| 1045 | """Construct a chunks tuple that broadcasts many chunks tuples |
| 1046 | |
| 1047 | >>> a = ((5, 5),) |
| 1048 | >>> b = ((5, 5),) |
| 1049 | >>> broadcast_chunks(a, b) |
| 1050 | ((5, 5),) |
| 1051 | |
| 1052 | >>> a = ((10, 10, 10), (5, 5),) |
| 1053 | >>> b = ((5, 5),) |
| 1054 | >>> broadcast_chunks(a, b) |
| 1055 | ((10, 10, 10), (5, 5)) |
| 1056 | |
| 1057 | >>> a = ((10, 10, 10), (5, 5),) |
| 1058 | >>> b = ((1,), (5, 5),) |
| 1059 | >>> broadcast_chunks(a, b) |
| 1060 | ((10, 10, 10), (5, 5)) |
| 1061 | |
| 1062 | >>> a = ((10, 10, 10), (5, 5),) |
| 1063 | >>> b = ((3, 3,), (5, 5),) |
| 1064 | >>> broadcast_chunks(a, b) |
| 1065 | Traceback (most recent call last): |
| 1066 | ... |
| 1067 | ValueError: Chunks do not align: [(10, 10, 10), (3, 3)] |
| 1068 | """ |
| 1069 | if not chunkss: |
| 1070 | return () |
| 1071 | elif len(chunkss) == 1: |
| 1072 | return chunkss[0] |
| 1073 | n = max(map(len, chunkss)) |
| 1074 | chunkss2 = [((1,),) * (n - len(c)) + c for c in chunkss] |
| 1075 | result = [] |
| 1076 | for i in range(n): |
| 1077 | step1 = [c[i] for c in chunkss2] |
| 1078 | if all(c == (1,) for c in step1): |
| 1079 | step2 = step1 |
| 1080 | else: |
| 1081 | step2 = [c for c in step1 if c != (1,)] |
| 1082 | if len(set(step2)) != 1: |
| 1083 | raise ValueError(f"Chunks do not align: {step2}") |
| 1084 | result.append(step2[0]) |
| 1085 | return tuple(result) |
| 1086 | |
| 1087 | |
| 1088 | def store( |
searching dependent graphs…