@@ -67,7 +67,12 @@ def track_num_layers(self):
67
67
track_num_tasks .unit = "tasks" # type: ignore[attr-defined] # Lazy
68
68
track_num_tasks_optimized .unit = "tasks" # type: ignore[attr-defined] # Lazy
69
69
track_num_layers .unit = "layers" # type: ignore[attr-defined] # Lazy
70
- for f in [track_num_tasks , track_num_tasks_optimized , track_num_layers , track_num_cohorts ]:
70
+ for f in [
71
+ track_num_tasks ,
72
+ track_num_tasks_optimized ,
73
+ track_num_layers ,
74
+ track_num_cohorts ,
75
+ ]:
71
76
f .repeat = 1 # type: ignore[attr-defined] # Lazy
72
77
f .rounds = 1 # type: ignore[attr-defined] # Lazy
73
78
f .number = 1 # type: ignore[attr-defined] # Lazy
@@ -82,9 +87,7 @@ def setup(self, *args, **kwargs):
82
87
y = np .repeat (np .arange (30 ), 60 )
83
88
by = x [np .newaxis , :] * y [:, np .newaxis ]
84
89
85
- self .by = flox .core ._factorize_multiple ((by ,), expected_groups = (None ,), any_by_dask = False )[
86
- 0
87
- ][0 ]
90
+ self .by = flox .core ._factorize_multiple ((by ,), expected_groups = (None ,), any_by_dask = False )[0 ][0 ]
88
91
89
92
self .array = dask .array .ones (self .by .shape , chunks = (350 , 350 ))
90
93
self .axis = (- 2 , - 1 )
@@ -101,7 +104,12 @@ def __init__(self, *args, **kwargs):
101
104
102
105
def rechunk (self ):
103
106
self .array = flox .core .rechunk_for_cohorts (
104
- self .array , - 1 , self .by , force_new_chunk_at = [1 ], chunksize = 48 , ignore_old_chunks = True
107
+ self .array ,
108
+ - 1 ,
109
+ self .by ,
110
+ force_new_chunk_at = [1 ],
111
+ chunksize = 48 ,
112
+ ignore_old_chunks = True ,
105
113
)
106
114
107
115
@@ -151,7 +159,12 @@ def setup(self, *args, **kwargs):
151
159
152
160
def rechunk (self ):
153
161
self .array = flox .core .rechunk_for_cohorts (
154
- self .array , - 1 , self .by , force_new_chunk_at = [1 ], chunksize = 4 , ignore_old_chunks = True
162
+ self .array ,
163
+ - 1 ,
164
+ self .by ,
165
+ force_new_chunk_at = [1 ],
166
+ chunksize = 4 ,
167
+ ignore_old_chunks = True ,
155
168
)
156
169
157
170
0 commit comments