Differential Evolution for nddimensionnal tensors

Hi all!

I tried using the Differential Evolution of tensorflow porbability, the example given in the tutorial works fine. However, in this tutorial the entries are 1 dimensionnal. So i tried using it for an objective function with a 2-d entries. i've taken the initial position as tf.constant([0.5,0.5]). Doing this i got the following error:
ValueError: Dimension 1 in both shapes must be equal, but are 2 and 1. Shapes are [50,2] and [50,1]. for 'Select_3' (op: 'Select') with input shapes: [50,1], [50,2], [50,2]. (50 is the population size)
The error comes from the one_step function in the update of the population:

to_replace = candidate_values < population_values
    next_population = [
        tf.compat.v1.where(to_replace, x=candidates_part, y=population_part)
        for candidates_part, population_part in zip(candidates, population)
    ]

to_replace is a [population_size,1] tensor while candidates_part and population_part are [population_size,population_dimension] tensors. And tf.where requires that the three tensors have the same shapes. So to avoid this problem i simply duplicate to_replace using to_replace = tf.tile(to_replace,[1,population_dimension])

Start at the Differential Evolution one_step population update shown in the issue, focusing on the shapes of to_replace, candidates_part, and population_part passed to tf.compat.v1.where. Reproduce the tutorial with a two-dimensional objective and verify that the population update completes without a shape error.