dpct.tensor.where wrong results with different shapes

[vlad-perevezentsev]

If shape of condition parameter is not equal to one of shapes of x or y , dpctl.tensor.where returns an incorrect result.
For example
cond_shape = (4,) , x_shape = (4,), y_shape=(4,) - correct
cond_shape = (2, 3) , x_shape = (2, 3), y_shape=(3,) - correct
cond_shape = (4,) , x_shape = (2, 3), y_shape=(2, 3) - incorrect

#NUMPY
np_cond = [False, True, True, False]
np_x = np.ones((3,4))
np_y = np.zeros((3,4))

#DPT 
dpt_cond = dpt.asarray(np_cond)
dpt_x = dpt.asarray(np_x)
dpt_y = dpt.asarray(np_y)

np.where(np_cond, np_x, np_y)
>> array([[0., 1., 1., 0.],
          [0., 1., 1., 0.],
          [0., 1., 1., 0.]])

dpt.where(dpt_cond, dpt_x, dpt_y)
>> usm_ndarray([[0., 0., 0., 1.],
                [1., 1., 1., 1.],
                [1., 0., 0., 0.]], dtype=float32)

[oleksandr-pavlyk]

Yes:

In [13]: dpt.where(dpt.broadcast_to(dpt_cond, (3,4)), dpt_x, dpt_y)
Out[13]:
usm_ndarray([[0., 0., 0., 1.],
             [1., 1., 1., 1.],
             [1., 0., 0., 0.]], dtype=float32)

In [14]: dpt.where(dpt.asarray(dpt.broadcast_to(dpt_cond, (3,4)), order='C', copy=True), dpt_x, dpt_y)
Out[14]:
usm_ndarray([[0., 1., 1., 0.],
             [0., 1., 1., 0.],
             [0., 1., 1., 0.]], dtype=float32)

The bug is in the strided kernel:

In [10]: with dpctl.utils.onetrace_enabled(): dpt.where(dpt.asarray(dpt.broadcast_to(dpt_cond, (3,4)), order='C', copy=True), dpt_x, dpt_y)
Device Timeline (queue: 0x55f8a9245a80): zeCommandListAppendMemoryCopy(H2D)<4.1> [ns] = 41159458105 (append) 41159880791 (submit) 41160644176 (start) 41160654801 (end)
Device Timeline (queue: 0x55f8a9245a80): dpctl::tensor::kernels::copy_and_cast::copy_cast_generic_kernel<bool, bool, dpctl::tensor::offset_utils::TwoOffsets_StridedIndexer><5.1> [ns] = 41227400991 (append) 41228018819 (submit) 41228987100 (start) 41229000954 (end)
Device Timeline (queue: 0x55f8a9245a80): dpctl::tensor::kernels::search::where_contig_kernel<float, bool, 4, 2><6.1> [ns] = 41284700850 (append) 41285005909 (submit) 41286470179 (start) 41286480283 (end)

In [11]: with dpctl.utils.onetrace_enabled(): dpt.where(dpt.broadcast_to(dpt_cond, (3,4)), dpt_x, dpt_y)
Device Timeline (queue: 0x55f8a9245a80): zeCommandListAppendMemoryCopy(H2D)<7.1> [ns] = 57067906835 (append) 57068241099 (submit) 57068562244 (start) 57068572973 (end)
Device Timeline (queue: 0x55f8a9245a80): dpctl::tensor::kernels::search::where_strided_kernel<float, bool, dpctl::tensor::offset_utils::ThreeOffsets_StridedIndexer><8.1> [ns] = 57130635631 (append) 57131152709 (submit) 57131977188 (start) 57132023750 (end)