Commit 6f19b45d authored by Theo Steininger's avatar Theo Steininger
Browse files

bincount in _slicing_distributor now tries to preserve a global_type...

bincount in _slicing_distributor now tries to preserve a global_type distribution strategy if possible.
parent 84f117d1
Pipeline #9803 failed with stage
in 9 minutes and 58 seconds
......@@ -1656,13 +1656,35 @@ class _slicing_distributor(distributor):
result_object = obj.copy_empty(global_shape=global_counts.shape,
result_object.set_local_data(global_counts, copy=False)
global_counts = local_counts
result_object = obj.copy_empty(local_shape=global_counts.shape,
global_shape = list(global_counts.shape)
global_shape[0] = self.global_shape[0]
# Try to preserve the distribution_strategy if possible.
# Choose the local distribution_strategy if...
# -> if the distribution_strategy was local in the beginning
# -> if the slicing of the first axis changes even though
# the bincounting wasn't performed on it. This can be the case
# if axis=0 is extremely small and collapsing axis=1,2,3...
# forces fftw to use a different slicing.
# Both cases are triggered by an ValueError exception
if self.distribution_strategy not in STRATEGIES['global']:
raise ValueError
result_object = obj.copy_empty(global_shape=global_shape,
result_object.set_local_data(global_counts, copy=False)
except ValueError:
result_object = obj.copy_empty(
result_object.set_local_data(global_counts, copy=False)
return result_object
......@@ -20,4 +20,4 @@
# 1) we don't load dependencies by storing it in
# 2) we can import it in for the same reason
# 3) we can import it into your module module
__version__ = '1.0.6'
__version__ = '1.0.7'
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment