diff --git a/HISTORY.rst b/HISTORY.rst index ba22ba79..1a04b858 100644 --- a/HISTORY.rst +++ b/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`) diff --git a/africanus/calibration/utils/dask.py b/africanus/calibration/utils/dask.py index 4c2655d1..073d84ff 100644 --- a/africanus/calibration/utils/dask.py +++ b/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)