# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import unittest import numpy as np from op_test import ( OpTest, convert_float_to_uint16, convert_uint16_to_float, get_device_place, is_custom_device, ) from utils import dygraph_guard, static_guard import paddle from paddle import base from paddle.autograd.ir_backward import grad from paddle.base import Program, core, program_guard from paddle.base.backward import append_backward class TestWhereOp(OpTest): def setUp(self): self.op_type = 'where' self.prim_op_type = 'prim' self.python_api = paddle.where self.public_python_api = paddle.where self.check_cinn = True self.init_config() self.inputs = {'Condition': self.cond, 'X': self.x, 'Y': self.y} self.outputs = {'Out': np.where(self.cond, self.x, self.y)} def test_check_output(self): self.check_output(check_cinn=self.check_cinn, check_pir=True) def test_check_grad(self): self.check_grad( ['X', 'Y'], 'Out', check_cinn=self.check_cinn, check_pir=True, check_prim_pir=True, ) def init_config(self): self.x = np.random.uniform((-3), 5, 100).astype('float64') self.y = np.random.uniform((-3), 5, 100).astype('float64') self.cond = np.zeros(100).astype('bool') class TestWhereOp2(TestWhereOp): def init_config(self): self.x = np.random.uniform((-5), 5, (60, 2)).astype('float64') self.y = np.random.uniform((-5), 5, (60, 2)).astype('float64') self.cond = np.ones((60, 2)).astype('bool') class TestWhereFP16OP(TestWhereOp): def init_config(self): self.dtype = np.float16 self.x = np.random.uniform((-5), 5, (60, 2)).astype(self.dtype) self.y = np.random.uniform((-5), 5, (60, 2)).astype(self.dtype) self.cond = np.ones((60, 2)).astype('bool') @unittest.skipIf( core.is_compiled_with_xpu(), "Skip complex due to lack of mean support", ) class TestWhereOpComplex64(TestWhereOp): def init_config(self): self.dtype = np.complex64 self.x = np.random.uniform((-5), 5, (60, 2)).astype(self.dtype) self.y = np.random.uniform((-5), 5, (60, 2)).astype(self.dtype) self.cond = np.ones((60, 2)).astype('bool') @unittest.skipIf( core.is_compiled_with_xpu(), "Skip complex due to lack of mean support", ) class TestWhereOpComplex128(TestWhereOp): def init_config(self): self.dtype = np.complex128 self.x = np.random.uniform((-5), 5, (60, 2)).astype(self.dtype) self.y = np.random.uniform((-5), 5, (60, 2)).astype(self.dtype) self.cond = np.ones((60, 2)).astype('bool') @unittest.skipIf( not (core.is_compiled_with_cuda() or is_custom_device()) or not core.is_bfloat16_supported(get_device_place()), "core is not compiled with CUDA and not support the bfloat16", ) class TestWhereBF16OP(OpTest): def setUp(self): self.op_type = 'where' self.prim_op_type = 'prim' self.dtype = np.uint16 self.python_api = paddle.where self.public_python_api = paddle.where self.check_cinn = True self.init_config() self.inputs = { 'Condition': self.cond, 'X': convert_float_to_uint16(self.x), 'Y': convert_float_to_uint16(self.y), } self.outputs = { 'Out': convert_float_to_uint16(np.where(self.cond, self.x, self.y)) } def test_check_output(self): place = get_device_place() self.check_output_with_place( place, check_cinn=self.check_cinn, check_pir=True ) def test_check_grad(self): place = get_device_place() self.check_grad_with_place( place, ['X', 'Y'], 'Out', numeric_grad_delta=0.05, check_cinn=self.check_cinn, check_pir=True, check_prim_pir=True, ) def init_config(self): self.x = np.random.uniform((-5), 5, (60, 2)).astype(np.float32) self.y = np.random.uniform((-5), 5, (60, 2)).astype(np.float32) self.cond = np.random.randint(2, size=(60, 2)).astype('bool') class TestWhereOp3(TestWhereOp): def init_config(self): self.x = np.random.uniform((-3), 5, (20, 2, 4)).astype('float64') self.y = np.random.uniform((-3), 5, (20, 2, 4)).astype('float64') self.cond = np.array(np.random.randint(2, size=(20, 2, 4)), dtype=bool) class TestWhereOp_ZeroSize(TestWhereOp): def init_config(self): self.x = np.random.uniform((-5), 5, (60, 0)).astype('float64') self.y = np.random.uniform((-5), 5, (60, 0)).astype('float64') self.cond = np.ones((60, 0)).astype('bool') class TestWhereAPI(unittest.TestCase): def setUp(self): self.init_data() def init_data(self): self.shape = [10, 15] self.cond = np.array(np.random.randint(2, size=self.shape), dtype=bool) self.x = np.random.uniform((-2), 3, self.shape).astype(np.float32) self.y = np.random.uniform((-2), 3, self.shape).astype(np.float32) self.out = np.where(self.cond, self.x, self.y) def ref_x_backward(self, dout): return np.where(self.cond, dout, 0) def ref_y_backward(self, dout): return np.where(~self.cond, dout, 0) def test_api(self, use_cuda=False): paddle.enable_static() for x_stop_gradient in [False, True]: for y_stop_gradient in [False, True]: with paddle.static.program_guard( paddle.static.Program(), paddle.static.Program() ): cond = paddle.static.data( name='cond', shape=[-1, *self.shape], dtype='bool' ) if not paddle.framework.use_pir_api(): cond.desc.set_need_check_feed(False) x = paddle.static.data( name='x', shape=[-1, *self.shape], dtype='float32' ) if not paddle.framework.use_pir_api(): x.desc.set_need_check_feed(False) y = paddle.static.data( name='y', shape=[-1, *self.shape], dtype='float32' ) if not paddle.framework.use_pir_api(): y.desc.set_need_check_feed(False) x.stop_gradient = x_stop_gradient if not paddle.framework.use_pir_api(): x.desc.set_need_check_feed(False) y.stop_gradient = y_stop_gradient if not paddle.framework.use_pir_api(): y.desc.set_need_check_feed(False) result = paddle.where(cond, x, y) result.stop_gradient = False append_backward(paddle.mean(result)) for use_cuda in [False, True]: if use_cuda and ( not ( base.core.is_compiled_with_cuda() or is_custom_device() ) ): break place = ( get_device_place() if use_cuda else base.CPUPlace() ) exe = base.Executor(place) if paddle.framework.use_pir_api(): fetch_list = [result] out = exe.run( paddle.static.default_main_program(), feed={ 'cond': self.cond, 'x': self.x, 'y': self.y, }, fetch_list=fetch_list, ) np.testing.assert_array_equal(out[0], self.out) else: fetch_list = [result, result.grad_name] if x_stop_gradient is False: fetch_list.append(x.grad_name) if y_stop_gradient is False: fetch_list.append(y.grad_name) out = exe.run( paddle.static.default_main_program(), feed={ 'cond': self.cond, 'x': self.x, 'y': self.y, }, fetch_list=fetch_list, ) np.testing.assert_array_equal(out[0], self.out) if x_stop_gradient is False: np.testing.assert_array_equal( out[2], self.ref_x_backward(out[1]) ) if y.stop_gradient is False: np.testing.assert_array_equal( out[3], self.ref_y_backward(out[1]) ) elif y.stop_gradient is False: np.testing.assert_array_equal( out[2], self.ref_y_backward(out[1]) ) paddle.disable_static() def test_pir_api(self, use_cuda=False): for x_stop_gradient in [False, True]: for y_stop_gradient in [False, True]: with ( paddle.pir_utils.IrGuard(), paddle.static.program_guard( paddle.static.Program(), paddle.static.Program() ), ): cond = paddle.static.data( name='cond', shape=self.shape, dtype='bool' ) x = paddle.static.data( name='x', shape=self.shape, dtype='float32' ) y = paddle.static.data( name='y', shape=self.shape, dtype='float32' ) x.stop_gradient = x_stop_gradient y.stop_gradient = y_stop_gradient result = paddle.where(cond, x, y) result.stop_gradient = False loss = paddle.mean(result) [x_grad, y_grad] = grad(loss, (x, y)) default_main_program = paddle.static.default_main_program() fetch_list = [result] if x_stop_gradient is False: fetch_list.append(x_grad) if y_stop_gradient is False: fetch_list.append(y_grad) for use_cuda in [False, True]: if use_cuda and ( not ( base.core.is_compiled_with_cuda() or is_custom_device() ) ): break place = ( get_device_place() if use_cuda else base.CPUPlace() ) exe = base.Executor(place) out = exe.run( default_main_program, feed={'cond': self.cond, 'x': self.x, 'y': self.y}, fetch_list=fetch_list, ) np.testing.assert_array_equal(out[0], self.out) if x_stop_gradient is False: np.testing.assert_array_equal( out[1], self.ref_x_backward(out[1]) ) if y.stop_gradient is False: np.testing.assert_array_equal( out[2], self.ref_y_backward(out[2]) ) elif y.stop_gradient is False: np.testing.assert_array_equal( out[1], self.ref_y_backward(out[1]) ) def test_api_broadcast(self, use_cuda=False): paddle.enable_static() main_program = paddle.static.Program() with paddle.static.program_guard(main_program): x = paddle.static.data(name='x', shape=[-1, 4, 1], dtype='float32') y = paddle.static.data(name='y', shape=[-1, 4, 2], dtype='float32') x_i = ( np.array([[0.9383, 0.1983, 3.2, 1.2]]) .astype('float32') .reshape([1, 4, 1]) ) y_i = ( np.array([[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]) .astype('float32') .reshape([1, 4, 2]) ) result = paddle.where((x > 1), x=x, y=y) for use_cuda in [False, True]: if use_cuda and ( not ( base.core.is_compiled_with_cuda() or is_custom_device() ) ): return place = get_device_place() if use_cuda else base.CPUPlace() exe = base.Executor(place) out = exe.run( paddle.static.default_main_program(), feed={'x': x_i, 'y': y_i}, fetch_list=[result], ) np.testing.assert_array_equal( out[0], np.where((x_i > 1), x_i, y_i) ) paddle.disable_static() def test_scalar(self): paddle.enable_static() main_program = paddle.static.Program() with paddle.static.program_guard(main_program): cond_shape = [4] cond = paddle.static.data( name='cond', shape=cond_shape, dtype='bool' ) x_data = 1.0 y_data = 2.0 cond_data = np.array([False, False, True, True]).astype('bool') result = paddle.where(condition=cond, x=x_data, y=y_data) for use_cuda in [False, True]: if use_cuda and ( not ( base.core.is_compiled_with_cuda() or is_custom_device() ) ): return place = get_device_place() if use_cuda else base.CPUPlace() exe = base.Executor(place) out = exe.run( paddle.static.default_main_program(), feed={'cond': cond_data}, fetch_list=[result], ) expect = np.where(cond_data, x_data, y_data) np.testing.assert_array_equal(out[0], expect) paddle.disable_static() def __test_where_with_broadcast_static(self, cond_shape, x_shape, y_shape): paddle.enable_static() main_program = paddle.static.Program() with paddle.static.program_guard(main_program): cond = paddle.static.data( name='cond', shape=cond_shape, dtype='bool' ) x = paddle.static.data(name='x', shape=x_shape, dtype='float32') y = paddle.static.data(name='y', shape=y_shape, dtype='float32') cond_data_tmp = np.random.random(size=cond_shape).astype('float32') cond_data = cond_data_tmp < 0.3 x_data = np.random.random(size=x_shape).astype('float32') y_data = np.random.random(size=y_shape).astype('float32') result = paddle.where(condition=cond, x=x, y=y) for use_cuda in [False, True]: if use_cuda and ( not ( base.core.is_compiled_with_cuda() or is_custom_device() ) ): return place = get_device_place() if use_cuda else base.CPUPlace() exe = base.Executor(place) out = exe.run( paddle.static.default_main_program(), feed={'cond': cond_data, 'x': x_data, 'y': y_data}, fetch_list=[result], ) expect = np.where(cond_data, x_data, y_data) np.testing.assert_array_equal(out[0], expect) def __test_where_with_type_promotion( self, x_dtype, y_dtype, expected_dtype=None ): paddle.enable_static() main_program = paddle.static.Program() shape = [3, 10] with paddle.static.program_guard(main_program): cond = paddle.static.data(name='cond', shape=[3, 10], dtype='bool') x = paddle.static.data(name='x', shape=shape, dtype=x_dtype) y = paddle.static.data(name='y', shape=shape, dtype=y_dtype) cond_data_tmp = np.random.random(size=shape).astype('float32') cond_data = cond_data_tmp < 0.3 if x_dtype != 'bfloat16': x_data = np.random.random(size=shape).astype(x_dtype) else: x_data = convert_float_to_uint16( np.random.random(size=shape).astype('float32') ) if y_dtype != 'bfloat16': y_data = np.random.random(size=shape).astype(y_dtype) else: y_data = convert_float_to_uint16( np.random.random(size=shape).astype('float32') ) result = paddle.where(condition=cond, x=x, y=y) for use_cuda in [False, True]: if use_cuda and ( not ( base.core.is_compiled_with_cuda() or is_custom_device() ) ): return place = get_device_place() if use_cuda else base.CPUPlace() exe = base.Executor(place) out = exe.run( paddle.static.default_main_program(), feed={'cond': cond_data, 'x': x_data, 'y': y_data}, fetch_list=[result], ) if x_dtype == 'bfloat16' or y_dtype == 'bfloat16': x_data_convert = ( convert_uint16_to_float(x_data) if x_dtype == 'bfloat16' else x_data ) y_data_convert = ( convert_uint16_to_float(y_data) if y_dtype == 'bfloat16' else y_data ) expect = np.where(cond_data, x_data_convert, y_data_convert) np.testing.assert_array_equal(out[0], expect) self.assertEqual(out[0].dtype.__str__(), expected_dtype) else: expect = np.where(cond_data, x_data, y_data) np.testing.assert_array_equal(out[0], expect) self.assertEqual(out[0].dtype, expect.dtype) def test_static_api_broadcast_1(self): cond_shape = [2, 4] a_shape = [2, 2, 4] b_shape = [2, 2, 4] self.__test_where_with_broadcast_static(cond_shape, a_shape, b_shape) def test_static_api_broadcast_2(self): cond_shape = [2, 1] a_shape = [2, 2, 4] b_shape = [2, 2, 4] self.__test_where_with_broadcast_static(cond_shape, a_shape, b_shape) def test_static_api_broadcast_3(self): cond_shape = [2, 2, 1] a_shape = [2, 2, 4] b_shape = [2, 2, 4] self.__test_where_with_broadcast_static(cond_shape, a_shape, b_shape) def test_static_api_broadcast_4(self): cond_shape = [2, 1, 4] a_shape = [2, 2, 4] b_shape = [2, 2, 4] self.__test_where_with_broadcast_static(cond_shape, a_shape, b_shape) def test_static_api_broadcast_5(self): cond_shape = [3, 2, 2, 4] a_shape = [2, 2, 4] b_shape = [2, 2, 4] self.__test_where_with_broadcast_static(cond_shape, a_shape, b_shape) def test_static_api_broadcast_6(self): cond_shape = [2, 2, 4] a_shape = [2, 2, 1] b_shape = [2, 2, 1] self.__test_where_with_broadcast_static(cond_shape, a_shape, b_shape) def test_static_api_broadcast_7(self): cond_shape = [2, 2, 4] a_shape = [2, 1, 4] b_shape = [2, 1, 4] self.__test_where_with_broadcast_static(cond_shape, a_shape, b_shape) def test_static_api_broadcast_8(self): cond_shape = [3, 2, 2, 4] a_shape = [2, 2, 1] b_shape = [2, 2, 1] self.__test_where_with_broadcast_static(cond_shape, a_shape, b_shape) def test_static_api_type_promotion_fp16_fp32(self): x_dtype = 'float16' y_dtype = 'float32' self.__test_where_with_type_promotion(x_dtype, y_dtype) self.__test_where_with_type_promotion(y_dtype, x_dtype) def test_static_api_type_promotion_fp16_fp64(self): x_dtype = 'float16' y_dtype = 'float64' self.__test_where_with_type_promotion(x_dtype, y_dtype) self.__test_where_with_type_promotion(y_dtype, x_dtype) def test_static_api_type_promotion_fp32_fp64(self): x_dtype = 'float32' y_dtype = 'float64' self.__test_where_with_type_promotion(x_dtype, y_dtype) self.__test_where_with_type_promotion(y_dtype, x_dtype) @unittest.skipIf( not ( (paddle.is_compiled_with_cuda() or is_custom_device()) and paddle.base.core.is_bfloat16_supported(get_device_place()) ), "bf16 is not supported in current device", ) def test_static_api_type_promotion_bf16_fp16(self): x_dtype = 'bfloat16' y_dtype = 'float16' self.__test_where_with_type_promotion(x_dtype, y_dtype, 'float32') self.__test_where_with_type_promotion(y_dtype, x_dtype, 'float32') @unittest.skipIf( not ( (paddle.is_compiled_with_cuda() or is_custom_device()) and paddle.base.core.is_bfloat16_supported(get_device_place()) ), "bf16 is not supported in current device", ) def test_static_api_type_promotion_bf16_fp32(self): x_dtype = 'bfloat16' y_dtype = 'float32' self.__test_where_with_type_promotion(x_dtype, y_dtype, 'float32') self.__test_where_with_type_promotion(y_dtype, x_dtype, 'float32') @unittest.skipIf( not ( (paddle.is_compiled_with_cuda() or is_custom_device()) and paddle.base.core.is_bfloat16_supported(get_device_place()) ), "bf16 is not supported in current device", ) def test_static_api_type_promotion_bf16_fp64(self): x_dtype = 'bfloat16' y_dtype = 'float64' self.__test_where_with_type_promotion(x_dtype, y_dtype, 'float64') self.__test_where_with_type_promotion(y_dtype, x_dtype, 'float64') class TestWhereDygraphAPI(unittest.TestCase): def test_api(self): with base.dygraph.guard(): x_i = np.array([0.9383, 0.1983, 3.2, 1.2]).astype('float64') y_i = np.array([1.0, 1.0, 1.0, 1.0]).astype('float64') cond_i = np.array([False, False, True, True]).astype('bool') x = paddle.to_tensor(x_i) y = paddle.to_tensor(y_i) cond = paddle.to_tensor(cond_i) out = paddle.where(cond, x, y) np.testing.assert_array_equal( out.numpy(), np.where(cond_i, x_i, y_i) ) def test_scalar(self): with base.dygraph.guard(): cond_i = np.array([False, False, True, True]).astype('bool') x = 1.0 y = 2.0 cond = paddle.to_tensor(cond_i) out = paddle.where(cond, x, y) np.testing.assert_array_equal(out.numpy(), np.where(cond_i, x, y)) def __test_where_with_broadcast_dygraph(self, cond_shape, a_shape, b_shape): with base.dygraph.guard(): cond_tmp = paddle.rand(cond_shape) cond = cond_tmp < 0.3 a = paddle.rand(a_shape) b = paddle.rand(b_shape) result = paddle.where(cond, a, b) result = result.numpy() expect = np.where(cond, a, b) np.testing.assert_array_equal(expect, result) def test_dygraph_api_broadcast_1(self): cond_shape = [2, 4] a_shape = [2, 2, 4] b_shape = [2, 2, 4] self.__test_where_with_broadcast_dygraph(cond_shape, a_shape, b_shape) def test_dygraph_api_broadcast_2(self): cond_shape = [2, 1] a_shape = [2, 2, 4] b_shape = [2, 2, 4] self.__test_where_with_broadcast_dygraph(cond_shape, a_shape, b_shape) def test_dygraph_api_broadcast_3(self): cond_shape = [2, 2, 1] a_shape = [2, 2, 4] b_shape = [2, 2, 4] self.__test_where_with_broadcast_dygraph(cond_shape, a_shape, b_shape) def test_dygraph_api_broadcast_4(self): cond_shape = [2, 1, 4] a_shape = [2, 2, 4] b_shape = [2, 2, 4] self.__test_where_with_broadcast_dygraph(cond_shape, a_shape, b_shape) def test_dygraph_api_broadcast_5(self): cond_shape = [3, 2, 2, 4] a_shape = [2, 2, 4] b_shape = [2, 2, 4] self.__test_where_with_broadcast_dygraph(cond_shape, a_shape, b_shape) def test_dygraph_api_broadcast_6(self): cond_shape = [2, 2, 4] a_shape = [2, 2, 1] b_shape = [2, 2, 1] self.__test_where_with_broadcast_dygraph(cond_shape, a_shape, b_shape) def test_dygraph_api_broadcast_7(self): cond_shape = [2, 2, 4] a_shape = [2, 1, 4] b_shape = [2, 1, 4] self.__test_where_with_broadcast_dygraph(cond_shape, a_shape, b_shape) def test_dygraph_api_broadcast_8(self): cond_shape = [3, 2, 2, 4] a_shape = [2, 2, 1] b_shape = [2, 2, 1] self.__test_where_with_broadcast_dygraph(cond_shape, a_shape, b_shape) def test_where_type_promotion_f2_f4(self): with base.dygraph.guard(): cond = np.array([False, False, True, True]).astype('bool') x = np.array([0.9383, 0.1983, 3.2, 1.2]).astype('float16') y = np.array([1.0, 1.0, 1.0, 1.0]).astype('float32') np.testing.assert_equal( paddle.float32, paddle.where( paddle.to_tensor(cond), paddle.to_tensor(x), paddle.to_tensor(y), ).dtype, ) np.testing.assert_array_equal( np.where(cond, x, y), paddle.where( paddle.to_tensor(cond), paddle.to_tensor(x), paddle.to_tensor(y), ).numpy(), ) np.testing.assert_array_equal( np.where(cond, y, x), paddle.where( paddle.to_tensor(cond), paddle.to_tensor(y), paddle.to_tensor(x), ).numpy(), ) def test_where_type_promotion_f2_f8(self): with base.dygraph.guard(): cond = np.array([False, False, True, True]).astype('bool') x = np.array([0.9383, 0.1983, 3.2, 1.2]).astype('float16') y = np.array([1.0, 1.0, 1.0, 1.0]).astype('float64') np.testing.assert_equal( paddle.float64, paddle.where( paddle.to_tensor(cond), paddle.to_tensor(x), paddle.to_tensor(y), ).dtype, ) np.testing.assert_array_equal( np.where(cond, x, y), paddle.where( paddle.to_tensor(cond), paddle.to_tensor(x), paddle.to_tensor(y), ).numpy(), ) np.testing.assert_array_equal( np.where(cond, y, x), paddle.where( paddle.to_tensor(cond), paddle.to_tensor(y), paddle.to_tensor(x), ).numpy(), ) def test_where_type_promotion_f4_f8(self): with base.dygraph.guard(): cond = np.array([False, False, True, True]).astype('bool') x = np.array([0.9383, 0.1983, 3.2, 1.2]).astype('float32') y = np.array([1.0, 1.0, 1.0, 1.0]).astype('float64') np.testing.assert_equal( paddle.float64, paddle.where( paddle.to_tensor(cond), paddle.to_tensor(x), paddle.to_tensor(y), ).dtype, ) np.testing.assert_array_equal( np.where(cond, x, y), paddle.where( paddle.to_tensor(cond), paddle.to_tensor(x), paddle.to_tensor(y), ).numpy(), ) np.testing.assert_array_equal( np.where(cond, y, x), paddle.where( paddle.to_tensor(cond), paddle.to_tensor(y), paddle.to_tensor(x), ).numpy(), ) def test_where_type_promotion_f2_bf(self): with base.dygraph.guard(): cond = np.array([False, False, True, True]).astype('bool') x = np.array([0.9383, 0.1983, 3.2, 1.2]).astype('float16') y = np.array([1.0, 1.0, 1.0, 1.0]).astype('float32') y = convert_uint16_to_float(convert_float_to_uint16(y)) np.testing.assert_equal( paddle.float32, paddle.where( paddle.to_tensor(cond), paddle.to_tensor(x), paddle.to_tensor(y), ).dtype, ) np.testing.assert_array_equal( np.where(cond, x, y), paddle.where( paddle.to_tensor(cond), paddle.to_tensor(x), paddle.to_tensor(y), ).numpy(), ) np.testing.assert_array_equal( np.where(cond, y, x), paddle.where( paddle.to_tensor(cond), paddle.to_tensor(y), paddle.to_tensor(x), ).numpy(), ) def test_where_type_promotion_f4_bf(self): with base.dygraph.guard(): cond = np.array([False, False, True, True]).astype('bool') x = np.array([0.9383, 0.1983, 3.2, 1.2]).astype('float32') y = np.array([1.0, 1.0, 1.0, 1.0]).astype('float32') y = convert_uint16_to_float(convert_float_to_uint16(y)) np.testing.assert_equal( paddle.float32, paddle.where( paddle.to_tensor(cond), paddle.to_tensor(x), paddle.to_tensor(y), ).dtype, ) np.testing.assert_array_equal( np.where(cond, x, y), paddle.where( paddle.to_tensor(cond), paddle.to_tensor(x), paddle.to_tensor(y), ).numpy(), ) np.testing.assert_array_equal( np.where(cond, y, x), paddle.where( paddle.to_tensor(cond), paddle.to_tensor(y), paddle.to_tensor(x), ), ) def test_where_type_promotion_f8_bf(self): with base.dygraph.guard(): cond = np.array([False, False, True, True]).astype('bool') x = np.array([0.9383, 0.1983, 3.2, 1.2]).astype('float64') y = np.array([1.0, 1.0, 1.0, 1.0]).astype('float32') y = convert_uint16_to_float(convert_float_to_uint16(y)) np.testing.assert_equal( paddle.float64, paddle.where( paddle.to_tensor(cond), paddle.to_tensor(x), paddle.to_tensor(y), ).dtype, ) np.testing.assert_array_equal( np.where(cond, x, y), paddle.where( paddle.to_tensor(cond), paddle.to_tensor(x), paddle.to_tensor(y), ).numpy(), ) np.testing.assert_array_equal( np.where(cond, y, x), paddle.where( paddle.to_tensor(cond), paddle.to_tensor(y), paddle.to_tensor(x), ).numpy(), ) def test_where_condition(self): data = np.array([[True, False], [False, True]]) with program_guard(Program(), Program()): x = paddle.static.data(name='x', shape=[(-1), 2], dtype='bool') if not paddle.framework.use_pir_api(): x.desc.set_need_check_feed(False) y = paddle.where(x) self.assertEqual(type(y), tuple) self.assertEqual(len(y), 2) z = paddle.concat(list(y), axis=0) exe = base.Executor(base.CPUPlace()) (res,) = exe.run( feed={'x': data}, fetch_list=[z], return_numpy=False ) expect_out = np.array([0, 1, 0, 1]) np.testing.assert_allclose(expect_out, np.array(res), rtol=1e-05) data = np.array([True, True, False]) with program_guard(Program(), Program()): x = paddle.static.data(name='x', shape=[-1], dtype='bool') if not paddle.framework.use_pir_api(): x.desc.set_need_check_feed(False) y = paddle.where(x) self.assertEqual(type(y), tuple) self.assertEqual(len(y), 1) z = paddle.concat(list(y), axis=0) exe = base.Executor(base.CPUPlace()) (res,) = exe.run( feed={'x': data}, fetch_list=[z], return_numpy=False ) expect_out = np.array([0, 1]) np.testing.assert_allclose(expect_out, np.array(res), rtol=1e-05) class TestWhereDygraphAPIBroadcast(unittest.TestCase): def test_broadcast_scalar(self): with base.dygraph.guard(): x_i = np.random.randn(4, 5, 6).astype('float64') y_i = -1.0 cond_i = np.random.randn(1, 1, 6).astype('bool') x = paddle.to_tensor(x_i) y = paddle.to_tensor(y_i) cond = paddle.to_tensor(cond_i) out = paddle.where(cond, x, y) np.testing.assert_array_equal( out.numpy(), np.where(cond_i, x_i, y_i) ) def test_broadcast_to_x(self): with base.dygraph.guard(): x_i = np.random.randn(4, 5, 6).astype('float64') y_i = np.random.randn(1, 5, 6).astype('float64') cond_i = np.random.randn(1, 1, 6).astype('bool') x = paddle.to_tensor(x_i) y = paddle.to_tensor(y_i) cond = paddle.to_tensor(cond_i) out = paddle.where(cond, x, y) np.testing.assert_array_equal( out.numpy(), np.where(cond_i, x_i, y_i) ) def test_broadcast_to_y(self): with base.dygraph.guard(): x_i = np.random.randn(1, 5, 6).astype('float64') y_i = np.random.randn(4, 5, 6).astype('float64') cond_i = np.random.randn(1, 1, 6).astype('bool') x = paddle.to_tensor(x_i) y = paddle.to_tensor(y_i) cond = paddle.to_tensor(cond_i) out = paddle.where(cond, x, y) np.testing.assert_array_equal( out.numpy(), np.where(cond_i, x_i, y_i) ) def test_broadcast_to_cond(self): with base.dygraph.guard(): x_i = np.random.randn(1, 1, 6).astype('float64') y_i = np.random.randn(1, 5, 1).astype('float64') cond_i = np.random.randn(4, 5, 6).astype('bool') x = paddle.to_tensor(x_i) y = paddle.to_tensor(y_i) cond = paddle.to_tensor(cond_i) out = paddle.where(cond, x, y) np.testing.assert_array_equal( out.numpy(), np.where(cond_i, x_i, y_i) ) def test_can_not_broadcast(self): with base.dygraph.guard(): x_i = np.random.randn(1, 1, 6).astype('float64') y_i = np.random.randn(1, 5, 3).astype('float64') cond_i = np.random.randn(4, 5, 6).astype('bool') x = paddle.to_tensor(x_i) y = paddle.to_tensor(y_i) cond = paddle.to_tensor(cond_i) with self.assertRaises(ValueError): _ = paddle.where(cond, x, y) class TestWhereDygraphAPIDtypePromotion(unittest.TestCase): def test_dtype_auto_promotion_float(self): with base.dygraph.guard(): x_i = np.random.randn(4, 5, 6).astype('float32') y_i = np.random.randn(4, 5, 6).astype('float64') cond_i = np.random.randn(4, 5, 6).astype('bool') x = paddle.to_tensor(x_i) y = paddle.to_tensor(y_i) cond = paddle.to_tensor(cond_i) out = paddle.where(cond, x, y) self.assertEqual(out.dtype, y.dtype) np.testing.assert_array_equal( out.numpy(), np.where(cond_i, x_i, y_i) ) class TestWhereOpError(unittest.TestCase): def test_errors(self): with paddle.static.program_guard( paddle.static.Program(), paddle.static.Program() ): x_i = np.array([0.9383, 0.1983, 3.2, 1.2]).astype('float64') y_i = np.array([1.0, 1.0, 1.0, 1.0]).astype('float64') cond_i = np.array([False, False, True, True]).astype('bool') def test_Variable(): paddle.where(cond_i, x_i, y_i) self.assertRaises(TypeError, test_Variable) def test_Value(): with paddle.pir_utils.IrGuard(): paddle.where(cond_i, x_i, y_i) self.assertRaises(TypeError, test_Value) def test_type(): with paddle.pir_utils.OldIrGuard(): x = paddle.static.data( name='x', shape=[-1, 4], dtype='bool' ) if not paddle.framework.use_pir_api(): x.desc.set_need_check_feed(False) y = paddle.static.data( name='y', shape=[-1, 4], dtype='float16' ) if not paddle.framework.use_pir_api(): y.desc.set_need_check_feed(False) cond = paddle.static.data( name='cond', shape=[-1, 4], dtype='int32' ) if not paddle.framework.use_pir_api(): cond.desc.set_need_check_feed(False) paddle.where(cond, x, y) self.assertRaises(TypeError, test_type) def test_value_error(self): with base.dygraph.guard(): cond_shape = [2, 2, 4] cond_tmp = paddle.rand(cond_shape) cond = cond_tmp < 0.3 a = paddle.rand(cond_shape) self.assertRaises(ValueError, paddle.where, cond, a) class TestWhereDygraphAPINonBoolCondition(unittest.TestCase): def test_condition_with_wrong_dtype(self): with base.dygraph.guard(): cond = paddle.to_tensor([True, False]) for dtype in [ paddle.int64, paddle.int32, paddle.float32, paddle.float64, ]: cond_wrong_dtype = cond.to(dtype) with self.assertRaises(ValueError): paddle.where(cond_wrong_dtype, 1, 0) def test_condition_inplace_with_wrong_dtype(self): with base.dygraph.guard(): cond = paddle.to_tensor([True, False]) x = paddle.zeros_like(cond).astype("float32") for dtype in [ paddle.int64, paddle.int32, paddle.float32, paddle.float64, ]: cond_wrong_dtype = cond.to(dtype) with self.assertRaises(ValueError): x = x.where_(cond_wrong_dtype, x, x) @unittest.skipIf( core.is_compiled_with_xpu(), "Skip XPU for zero size tensor is not fully supported", ) class TestWhereZeroSizeTensor(unittest.TestCase): def init_inputs_outputs(self, shapes): cond = np.random.randint(0, 2, size=shapes[0]).astype('bool') x = np.random.random(shapes[1]).astype('float64') y = np.random.random(shapes[2]).astype('float64') out_ref = np.where(cond, x, y) return (cond, x, y), out_ref def _test_with_shapes(self, shapes): inputs, out_ref = self.init_inputs_outputs(shapes) with dygraph_guard(): tensors = [paddle.to_tensor(inp) for inp in inputs] result = paddle.where(tensors[0], tensors[1], tensors[2]) np.testing.assert_allclose(result, out_ref, rtol=1e-05) with ( static_guard(), paddle.static.program_guard( paddle.static.Program(), paddle.static.Program() ), ): cond_t = paddle.static.data( name='cond', shape=[2, 3, 5], dtype='bool' ) x_t = paddle.static.data(name='x', shape=[2, 3, 5], dtype='float64') y_t = paddle.static.data(name='y', shape=[2, 3, 5], dtype='float64') result = paddle.where(cond_t, x_t, y_t) exe = base.Executor(base.CPUPlace()) out = exe.run( paddle.static.default_main_program(), feed={'cond': inputs[0], 'x': inputs[1], 'y': inputs[2]}, fetch_list=[result], ) np.testing.assert_allclose(out[0], out_ref, rtol=1e-05) def test_api_with_zero_size_input(self): self._test_with_shapes([(0, 0), (0, 0), (0, 0)]) self._test_with_shapes([(0, 1), (0, 1), (0, 1)]) self._test_with_shapes([(0, 2, 1), (0, 2, 1), (0, 2, 1)]) self._test_with_shapes([(5, 17, 0, 6), (5, 17, 0, 6), (5, 17, 0, 6)]) self._test_with_shapes([(0, 5, 17, 6), (0, 5, 17, 6), (0, 5, 17, 6)]) @unittest.skipIf( core.is_compiled_with_xpu(), "Skip XPU for bool dtype is not fully supported", ) class TestWhereBoolInput(unittest.TestCase): def test_api_with_dygraph(self): cond = np.random.randint(0, 2, size=[2, 3, 5]).astype('bool') x = np.random.random([2, 3, 5]).astype('bool') y = np.random.random([2, 3, 5]).astype('bool') out_ref = np.where(cond, x, y) with dygraph_guard(): cond_t = paddle.to_tensor(cond) x_t = paddle.to_tensor(x) y_t = paddle.to_tensor(y) result = paddle.where(cond_t, x_t, y_t) np.testing.assert_allclose(result, out_ref, rtol=1e-05) def test_api_with_static(self): cond = np.random.randint(0, 2, size=[2, 3, 5]).astype('bool') x = np.random.random([2, 3, 5]).astype('bool') y = np.random.random([2, 3, 5]).astype('bool') out_ref = np.where(cond, x, y) with ( static_guard(), paddle.static.program_guard( paddle.static.Program(), paddle.static.Program() ), ): cond_t = paddle.static.data( name='cond', shape=[2, 3, 5], dtype='bool' ) x_t = paddle.static.data(name='x', shape=[2, 3, 5], dtype='bool') y_t = paddle.static.data(name='y', shape=[2, 3, 5], dtype='bool') result = paddle.where(cond_t, x_t, y_t) exe = base.Executor(base.CPUPlace()) out = exe.run( paddle.static.default_main_program(), feed={'cond': cond, 'x': x, 'y': y}, fetch_list=[result], ) np.testing.assert_allclose(out[0], out_ref, rtol=1e-05) class TestWhereAlias(unittest.TestCase): def setUp(self): paddle.disable_static() def test_where_alias(self): """ Test the alias of where function. ``where(condition=cond, input=x, other=y)`` is equivalent to ``where(condition=cond, x=x, y=y)`` """ shape = [2, 4] cond = paddle.randint(0, 2, shape).astype("bool") x = paddle.rand(shape).astype("float32") y = paddle.rand(shape).astype("float32") # Test all alias combinations combinations = [ {"condition": cond, "x": x, "y": y}, {"condition": cond, "input": x, "y": y}, {"condition": cond, "x": x, "other": y}, {"condition": cond, "input": x, "other": y}, ] # Get baseline result expected = np.where(cond.numpy(), x.numpy(), y.numpy()) for params in combinations: out = paddle.where(**params) np.testing.assert_allclose(out.numpy(), expected, rtol=1e-05) paddle.enable_static() class TestWhereOut(unittest.TestCase): def setUp(self): self.cond_np = np.random.randint(0, 2, size=[2, 3, 5]).astype('bool') self.x_np = np.random.random([2, 3, 5]).astype('float32') self.y_np = np.random.random([2, 3, 5]).astype('float32') def test_api_with_dygraph(self): paddle.disable_static() cond = paddle.to_tensor(self.cond_np) x = paddle.to_tensor(self.x_np) y = paddle.to_tensor(self.y_np) out_holder = paddle.zeros_like(cond) out_ref = paddle.where(cond, x, y) paddle.where(cond, x, y, out=out_holder) np.testing.assert_allclose(out_holder, out_ref, rtol=1e-20) if __name__ == "__main__": paddle.enable_static() unittest.main()