Add tests for `returned=True`
@@ -161,45 +161,34 @@ class TestAverage:
@testing.for_all_dtypes()
@testing.numpy_cupy_allclose()
- def test_average_axis_weights(self, xp, dtype):
- a = testing.shaped_arange((2, 3, 4), xp, dtype)
- w = testing.shaped_arange((2, 3, 4), xp, dtype)
- return xp.average(a, axis=2, weights=w)
-
- def check_returned(self, a, axis, weights):
- average_cpu, sum_weights_cpu = numpy.average(
- a, axis, weights, returned=True)
- result = cupy.average(
- cupy.asarray(a), axis, weights, returned=True)
- assert isinstance(result, tuple)
- assert len(result) == 2
- average_gpu, sum_weights_gpu = result
- testing.assert_allclose(average_cpu, average_gpu)
- testing.assert_allclose(sum_weights_cpu, sum_weights_gpu)
-
- @testing.for_all_dtypes()
- def test_returned(self, dtype):
+ @pytest.mark.parametrize(
+ 'axis,weights', [(1, False), (None, True), (1, True)])
+ def test_returned(self, xp, dtype, axis, weights):
a = testing.shaped_arange((2, 3), numpy, dtype)
+ if weights:
w = testing.shaped_arange((2, 3), numpy, dtype)
- self.check_returned(a, axis=1, weights=None)
- self.check_returned(a, axis=None, weights=w)
- self.check_returned(a, axis=1, weights=w)
+ else:
+ w = None
+ return xp.average(a, axis=axis, weights=w, returned=True)
@testing.for_all_dtypes()
@testing.numpy_cupy_allclose(rtol={'default': 5e-7})
+ @pytest.mark.parametrize('returned', [True, False])
@testing.with_requires('numpy>=1.23')
- def test_average_keepdims_axis1(self, xp, dtype):
+ def test_average_keepdims_axis1(self, xp, dtype, returned):
a = testing.shaped_random((2, 3), xp, dtype)
w = testing.shaped_random((2, 3), xp, dtype)
- return xp.average(a, axis=1, weights=w, keepdims=True)
+ return xp.average(
+ a, axis=1, weights=w, returned=returned, keepdims=True)
@testing.for_all_dtypes()
@testing.numpy_cupy_allclose(rtol={'default': 1e-7, numpy.float16: 1e-3})
+ @pytest.mark.parametrize('returned', [True, False])
@testing.with_requires('numpy>=1.23')
- def test_average_keepdims_noaxis(self, xp, dtype):
+ def test_average_keepdims_noaxis(self, xp, dtype, returned):
a = testing.shaped_random((2, 3), xp, dtype)
w = testing.shaped_random((2, 3), xp, dtype)
- return xp.average(a, weights=w, keepdims=True)
+ return xp.average(a, weights=w, returned=returned, keepdims=True)
@testing.gpu
