@@ -67,7 +67,12 @@ def track_num_layers(self):
6767 track_num_tasks .unit = "tasks" # type: ignore[attr-defined] # Lazy
6868 track_num_tasks_optimized .unit = "tasks" # type: ignore[attr-defined] # Lazy
6969 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+ ]:
7176 f .repeat = 1 # type: ignore[attr-defined] # Lazy
7277 f .rounds = 1 # type: ignore[attr-defined] # Lazy
7378 f .number = 1 # type: ignore[attr-defined] # Lazy
@@ -82,9 +87,7 @@ def setup(self, *args, **kwargs):
8287 y = np .repeat (np .arange (30 ), 60 )
8388 by = x [np .newaxis , :] * y [:, np .newaxis ]
8489
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 ]
8891
8992 self .array = dask .array .ones (self .by .shape , chunks = (350 , 350 ))
9093 self .axis = (- 2 , - 1 )
@@ -101,7 +104,12 @@ def __init__(self, *args, **kwargs):
101104
102105 def rechunk (self ):
103106 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 ,
105113 )
106114
107115
@@ -151,7 +159,12 @@ def setup(self, *args, **kwargs):
151159
152160 def rechunk (self ):
153161 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 ,
155168 )
156169
157170
0 commit comments