def test_select(): element_type = Type.f32 A = Parameter(Type.boolean, Shape([1, 2])) B = Parameter(element_type, Shape([1, 2])) C = Parameter(element_type, Shape([1, 2])) parameter_list = [A, B, C] function = Function([ng.select(A, B, C)], parameter_list, 'test') backend = Backend.create(test.BACKEND_NAME) a = backend.create_tensor(Type.boolean, Shape([1, 2])) b = backend.create_tensor(element_type, Shape([1, 2])) c = backend.create_tensor(element_type, Shape([1, 2])) result = backend.create_tensor(element_type, Shape([1, 2])) a.write(util.numpy_to_c(np.array([[True, False]], dtype=np.bool)), 2) b.write(util.numpy_to_c(np.array([[5, 6]], dtype=np.float32)), 8) c.write(util.numpy_to_c(np.array([[7, 8]], dtype=np.float32)), 8) result_arr = np.array([[0, 0]], dtype=np.float32) result.write(util.numpy_to_c(result_arr), 8) handle = backend.compile(function) handle.call([result], [a, b, c]) result.read(util.numpy_to_c(result_arr), 8) result_arr_ref = np.array([[5, 8]]) assert np.allclose(result_arr, result_arr_ref)
def test_select(): element_type = Type.f32 A = Parameter(Type.boolean, Shape([1, 2])) B = Parameter(element_type, Shape([1, 2])) C = Parameter(element_type, Shape([1, 2])) parameter_list = [A, B, C] function = Function([ng.select(A, B, C)], parameter_list, "test") runtime = get_runtime() computation = runtime.computation(function, *parameter_list) result = computation( np.array([[True, False]], dtype=np.bool), np.array([[5, 6]], dtype=np.float32), np.array([[7, 8]], dtype=np.float32), )[0] expected = np.array([[5, 8]]) assert np.allclose(result, expected)