@@ -614,52 +614,83 @@ def hybrid_forward(self, F, a, *args, **kwargs):
614614 def is_int (dtype ):
615615 return 'int' in dtype
616616
617+ is_windows = sys .platform .startswith ('win' )
617618 in_data_dim = random .choice ([2 , 3 , 4 ])
618619 shape = rand_shape_nd (in_data_dim , dim = 3 )
619620 acc_type = {'float16' : 'float32' , 'float32' : 'float64' , 'float64' : 'float64' ,
620- 'int8' : 'int32' , 'int32' : 'int64' , 'int64' : 'int64' }
621+ 'bool' : 'int64' , 'int8' : 'int32' , 'int32' : 'int64' , 'int64' : 'int64' }
622+ ft_types = ['float16' , 'float32' , 'float64' ]
623+ it_types = ['bool' , 'int8' , 'int32' , 'int64' ]
621624 for hybridize in [False , True ]:
622625 for keepdims in [True , False ]:
623626 for axis in ([i for i in range (in_data_dim )] + [(), None ]):
624- for itype in ['float16' , 'float32' , 'float64' ]:
625- for dtype in ['float16' , 'float32' , 'float64' ]:
626- if is_int (dtype ) and not is_int (itype ):
627- continue
628- # test gluon
629- test_mean = TestMean (axis = axis , dtype = dtype , keepdims = keepdims )
630- if hybridize :
631- test_mean .hybridize ()
632- if is_int (itype ):
633- x = _np .random .randint (- 128 , 128 , shape , dtype = itype )
634- x = mx .nd .array (x , dtype = itype )
635- else :
636- x = mx .nd .random .uniform (- 1.0 , 1.0 , shape = shape , dtype = itype )
637- x = x .as_np_ndarray ()
638- x .attach_grad ()
627+ for itype , dtype in itertools .product (ft_types , [None ] + ft_types + it_types ):
628+ if dtype == 'bool' :
629+ continue
630+ # test gluon
631+ test_mean = TestMean (axis = axis , dtype = dtype , keepdims = keepdims )
632+ if hybridize :
633+ test_mean .hybridize ()
634+ x = np .random .uniform (- 1.0 , 1.0 , size = shape ).astype (itype )
635+ x = x .as_np_ndarray ()
636+ x .attach_grad ()
639637
640- expected_ret = _np .mean (x .asnumpy (), axis = axis , dtype = acc_type [itype ], keepdims = keepdims )
641- expected_ret = expected_ret .astype (dtype )
642- with mx .autograd .record ():
643- y = test_mean (x )
644- assert y .shape == expected_ret .shape
645- assert_almost_equal (y .asnumpy (), expected_ret , rtol = 1e-3 if dtype == 'float16' else 1e-3 ,
646- atol = 1e-5 if dtype == 'float16' else 1e-5 )
638+ expected_ret = _np .mean (x .asnumpy (), axis = axis , dtype = acc_type [itype ], keepdims = keepdims )
639+ expected_ret = expected_ret .astype (dtype )
640+ with mx .autograd .record ():
641+ y = test_mean (x )
642+ assert y .shape == expected_ret .shape
643+ assert_almost_equal (y .asnumpy (), expected_ret , rtol = 1e-3 if dtype == 'float16' else 1e-3 ,
644+ atol = 1e-5 if dtype == 'float16' else 1e-5 )
647645
648- y .backward ()
649- N = x .size / y .size
650- assert same (x .grad .asnumpy (), _np .ones (shape = x .shape , dtype = x .dtype ) / N )
646+ y .backward ()
647+ N = x .size / y .size
648+ assert same (x .grad .asnumpy (), _np .ones (shape = x .shape , dtype = x .dtype ) / N )
651649
652- # test numeric
653- if itype == 'float32' and dtype == 'float32' :
654- x_sym = mx .sym .Variable ("x" ).as_np_ndarray ()
655- mx_sym = mx .sym .np .mean (x_sym , axis = axis , dtype = dtype , keepdims = keepdims ).as_nd_ndarray ()
656- check_numeric_gradient (mx_sym , [x .as_nd_ndarray ()],
657- numeric_eps = 1e-3 , rtol = 1e-3 , atol = 1e-4 , dtype = _np .float32 )
650+ # test numeric
651+ if itype == 'float32' and dtype == 'float32' :
652+ x_sym = mx .sym .Variable ("x" ).as_np_ndarray ()
653+ mx_sym = mx .sym .np .mean (x_sym , axis = axis , dtype = dtype , keepdims = keepdims ).as_nd_ndarray ()
654+ check_numeric_gradient (mx_sym , [x .as_nd_ndarray ()],
655+ numeric_eps = 1e-3 , rtol = 1e-3 , atol = 1e-4 , dtype = _np .float32 )
658656
659- # test imperative
660- mx_out = np .mean (x , axis = axis , dtype = dtype , keepdims = keepdims )
661- np_out = _np .mean (x .asnumpy (), axis = axis , dtype = acc_type [itype ], keepdims = keepdims ).astype (dtype )
662- assert_almost_equal (mx_out .asnumpy (), np_out , rtol = 1e-3 , atol = 1e-5 )
657+ # test imperative
658+ mx_out = np .mean (x , axis = axis , dtype = dtype , keepdims = keepdims )
659+ np_out = _np .mean (x .asnumpy (), axis = axis , dtype = acc_type [itype ], keepdims = keepdims ).astype (dtype )
660+ assert_almost_equal (mx_out .asnumpy (), np_out , rtol = 1e-3 , atol = 1e-5 )
661+
662+ for itype , dtype in itertools .product (it_types , [None ] + ft_types + it_types ):
663+ if dtype == 'bool' :
664+ continue
665+ # test gluon
666+ test_mean = TestMean (axis = axis , dtype = dtype , keepdims = keepdims )
667+ if hybridize :
668+ test_mean .hybridize ()
669+
670+ if itype == 'bool' :
671+ x = np .array (_np .random .uniform (size = shape ) > 0.5 )
672+ else :
673+ x = np .random .uniform (- 128 , 127 , size = shape ).astype (itype )
674+
675+ expected_ret = _np .mean (x .asnumpy (), axis = axis , dtype = dtype , keepdims = keepdims )
676+
677+ if itype == 'bool' :
678+ if is_op_runnable () and (not is_windows ) and dtype not in 'float16' , 'int8' ]: # special handling of boolean ndarray
679+ y = test_mean (x )
680+ assert y .shape == expected_ret .shape
681+ assert_almost_equal (y .asnumpy (), expected_ret , rtol = 1e-3 if dtype == 'float16' else 1e-3 ,
682+ atol = 1e-5 if dtype == 'float16' else 1e-5 )
683+ continue
684+
685+ y = test_mean (x )
686+ assert y .shape == expected_ret .shape
687+ assert_almost_equal (y .asnumpy (), expected_ret , rtol = 1e-3 if dtype == 'float16' else 1e-3 ,
688+ atol = 1e-5 if dtype == 'float16' else 1e-5 )
689+
690+ # test imperative
691+ mx_out = np .mean (x , axis = axis , dtype = dtype , keepdims = keepdims )
692+ np_out = _np .mean (x .asnumpy (), axis = axis , dtype = dtype , keepdims = keepdims ).astype (dtype )
693+ assert_almost_equal (mx_out .asnumpy (), np_out , rtol = 1e-3 , atol = 1e-5 )
663694
664695
665696@with_seed ()
0 commit comments