<s> import unittest <EOL> from cupy import testing <EOL> @ testing . gpu <EOL> class TestSearch ( unittest . TestCase ) : <EOL> _multiprocess_can_split_ = True <EOL> @ testing . for_all_dtypes ( ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_argmax_all ( self , xp , dtype ) : <EOL> a = testing . shaped_random ( ( <NUM_LIT:2> , <NUM_LIT:3> ) , xp , dtype ) <EOL> return a . argmax ( ) <EOL> @ testing . for_all_dtypes ( ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_external_argmax_all ( self , xp , dtype ) : <EOL> a = testing . shaped_random ( ( <NUM_LIT:2> , <NUM_LIT:3> ) , xp , dtype ) <EOL> return xp . argmax ( a ) <EOL> @ testing . for_all_dtypes ( ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_argmax_axis_large ( self , xp , dtype ) : <EOL> a = testing . shaped_random ( ( <NUM_LIT:3> , <NUM_LIT:1000> ) , xp , dtype ) <EOL> return a . argmax ( axis = <NUM_LIT:0> ) <EOL> @ testing . for_all_dtypes ( ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_external_argmax_axis_large ( self , xp , dtype ) : <EOL> a = testing . shaped_random ( ( <NUM_LIT:3> , <NUM_LIT:1000> ) , xp , dtype ) <EOL> return xp . argmax ( a , axis = <NUM_LIT:0> ) <EOL> @ testing . for_all_dtypes ( ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_argmax_axis0 ( self , xp , dtype ) : <EOL> a = testing . shaped_random ( ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) , xp , dtype ) <EOL> return a . argmax ( axis = <NUM_LIT:0> ) <EOL> @ testing . for_all_dtypes ( ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_argmax_axis1 ( self , xp , dtype ) : <EOL> a = testing . shaped_random ( ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) , xp , dtype ) <EOL> return a . argmax ( axis = <NUM_LIT:1> ) <EOL> @ testing . for_all_dtypes ( ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_argmax_axis2 ( self , xp , dtype ) : <EOL> a = testing . shaped_random ( ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) , xp , dtype ) <EOL> return a . argmax ( axis = <NUM_LIT:2> ) <EOL> @ testing . for_all_dtypes ( ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_argmin_all ( self , xp , dtype ) : <EOL> a = testing . shaped_random ( ( <NUM_LIT:2> , <NUM_LIT:3> ) , xp , dtype ) <EOL> return a . argmin ( ) <EOL> @ testing . for_all_dtypes ( ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_external_argmin_all ( self , xp , dtype ) : <EOL> a = testing . shaped_random ( ( <NUM_LIT:2> , <NUM_LIT:3> ) , xp , dtype ) <EOL> return xp . argmin ( a ) <EOL> @ testing . for_all_dtypes ( ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_argmin_axis_large ( self , xp , dtype ) : <EOL> a = testing . shaped_random ( ( <NUM_LIT:3> , <NUM_LIT:1000> ) , xp , dtype ) <EOL> return a . argmin ( axis = <NUM_LIT:0> ) <EOL> @ testing . for_all_dtypes ( ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_external_argmin_axis_large ( self , xp , dtype ) : <EOL> a = testing . shaped_random ( ( <NUM_LIT:3> , <NUM_LIT:1000> ) , xp , dtype ) <EOL> return xp . argmin ( a , axis = <NUM_LIT:0> ) <EOL> @ testing . for_all_dtypes ( ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_argmin_axis0 ( self , xp , dtype ) : <EOL> a = testing . shaped_random ( ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) , xp , dtype ) <EOL> return a . argmin ( axis = <NUM_LIT:0> ) <EOL> @ testing . for_all_dtypes ( ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_argmin_axis1 ( self , xp , dtype ) : <EOL> a = testing . shaped_random ( ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) , xp , dtype ) <EOL> return a . argmin ( axis = <NUM_LIT:1> ) <EOL> @ testing . for_all_dtypes ( ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_argmin_axis2 ( self , xp , dtype ) : <EOL> a = testing . shaped_random ( ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) , xp , dtype ) <EOL> return a . argmin ( axis = <NUM_LIT:2> ) <EOL> @ testing . parameterize ( <EOL> { '<STR_LIT>' : ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) , '<STR_LIT>' : ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) , '<STR_LIT>' : ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) } , <EOL> { '<STR_LIT>' : ( <NUM_LIT:4> , ) , '<STR_LIT>' : ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) , '<STR_LIT>' : ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) } , <EOL> { '<STR_LIT>' : ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) , '<STR_LIT>' : ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) , '<STR_LIT>' : ( <NUM_LIT:3> , <NUM_LIT:4> ) } , <EOL> { '<STR_LIT>' : ( <NUM_LIT:3> , <NUM_LIT:4> ) , '<STR_LIT>' : ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) , '<STR_LIT>' : ( <NUM_LIT:4> , ) } , <EOL> ) <EOL> @ testing . gpu <EOL> class TestWhere ( unittest . TestCase ) : <EOL> @ testing . for_all_dtypes_combination ( <EOL> names = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> @ testing . numpy_cupy_allclose ( ) <EOL> def test_where ( self , xp , cond_type , x_type , y_type ) : <EOL> m = testing . shaped_random ( self . cond_shape , xp , xp . bool_ ) <EOL> cond = testing . shaped_random ( self . cond_shape , xp , cond_type ) * m <EOL> x = testing . shaped_random ( self . x_shape , xp , x_type ) <EOL> y = testing . shaped_random ( self . y_shape , xp , y_type ) <EOL> return xp . where ( cond , x , y ) <EOL> @ testing . gpu <EOL> class TestWhereError ( unittest . TestCase ) : <EOL> @ testing . numpy_cupy_raises ( ) <EOL> def test_one_argument ( self , xp ) : <EOL> cond = testing . shaped_random ( ( <NUM_LIT:3> , <NUM_LIT:4> ) , xp , dtype = xp . bool_ ) <EOL> x = testing . shaped_random ( ( <NUM_LIT:2> , <NUM_LIT:3> , <NUM_LIT:4> ) , xp , xp . int32 ) <EOL> xp . where ( cond , x ) </s>
