remove new_axes={'corr2': 2} in calibration utils (#288)

HISTORY.rst

@@ -4,6 +4,7 @@ History
 
 X.Y.Z (YYYY-MM-DD)
 ------------------
+* Remove unnecessary new_axes in calibration utils after upstream fix in dask (:pr:`288`)
 * Check that ncorr is never larger than 2 in calibration utils (:pr:`287`)
 * Optionally check NRT allocations (:pr:`286`)
 * Use `packaging.version.Version` instead of deprecated `distutils.version.LooseVersion` (:pr:`285`)

africanus/calibration/utils/dask.py

@@ -56,8 +56,6 @@ def corrupt_vis(time_bin_indices, time_bin_counts, antenna1,
     else:
         raise ValueError("Unknown mode argument of %s" % mode)
 
-    # the new_axes={"corr2": 2} is required because of a dask bug
-    # see https://github.com/dask/dask/issues/5550
     return blockwise(_corrupt_vis_wrapper, out_shape,
                      time_bin_indices, ("row",),
                      time_bin_counts, ("row",),
@@ -66,7 +64,6 @@ def corrupt_vis(time_bin_indices, time_bin_counts, antenna1,
                      jones, jones_shape,
                      model, model_shape,
                      adjust_chunks={"row": antenna1.chunks[0]},
-                     new_axes={"corr2": 2},
                      dtype=model.dtype,
                      align_arrays=False)
 
@@ -114,8 +111,6 @@ def compute_and_corrupt_vis(time_bin_indices, time_bin_counts,
     else:
         raise ValueError("Unknown mode argument of %s" % mode)
 
-    # the new_axes={"corr2": 2} is required because of a dask bug
-    # see https://github.com/dask/dask/issues/5550
     return blockwise(_compute_and_corrupt_vis_wrapper, out_shape,
                      time_bin_indices, ("row",),
                      time_bin_counts, ("row",),
@@ -127,7 +122,6 @@ def compute_and_corrupt_vis(time_bin_indices, time_bin_counts,
                      freq, ("chan",),
                      lm, ("row", "dir", "two"),
                      adjust_chunks={"row": antenna1.chunks[0]},
-                     new_axes={"corr2": 2},
                      dtype=model.dtype,
                      align_arrays=False)
 
@@ -161,8 +155,6 @@ def correct_vis(time_bin_indices, time_bin_counts, antenna1,
     else:
         raise ValueError("Unknown mode argument of %s" % mode)
 
-    # the new_axes={"corr2": 2} is required because of a dask bug
-    # see https://github.com/dask/dask/issues/5550
     return blockwise(_correct_vis_wrapper, out_shape,
                      time_bin_indices, ("row",),
                      time_bin_counts, ("row",),
@@ -172,7 +164,6 @@ def correct_vis(time_bin_indices, time_bin_counts, antenna1,
                      vis, out_shape,
                      flag, out_shape,
                      adjust_chunks={"row": antenna1.chunks[0]},
-                     new_axes={"corr2": 2},
                      dtype=vis.dtype,
                      align_arrays=False)
 
@@ -211,8 +202,6 @@ def residual_vis(time_bin_indices, time_bin_counts, antenna1,
     else:
         raise ValueError("Unknown mode argument of %s" % mode)
 
-    # the new_axes={"corr2": 2} is required because of a dask bug
-    # see https://github.com/dask/dask/issues/5550
     return blockwise(_residual_vis_wrapper, out_shape,
                      time_bin_indices, ("row",),
                      time_bin_counts, ("row",),
@@ -223,7 +212,6 @@ def residual_vis(time_bin_indices, time_bin_counts, antenna1,
                      flag, out_shape,
                      model, model_shape,
                      adjust_chunks={"row": antenna1.chunks[0]},
-                     new_axes={"corr2": 2},
                      dtype=vis.dtype,
                      align_arrays=False)