def test_keepdims_wrapper_no_axis(): def summer(a, axis=None): return a.sum(axis=axis) summer_wrapped = keepdims_wrapper(summer) assert summer_wrapped != summer assert summer_wrapped == keepdims_wrapper(summer_wrapped) a = np.arange(24).reshape(1, 2, 3, 4) r = summer(a) rw = summer_wrapped(a, keepdims=True) rwf = summer_wrapped(a, keepdims=False) assert r.ndim == 0 assert r.shape == tuple() assert r == 276 assert rw.ndim == 4 assert rw.shape == (1, 1, 1, 1) assert (rw == 276).all() assert rwf.ndim == 0 assert rwf.shape == tuple() assert rwf == 276
def test_keepdims_wrapper_two_axes(): def summer(a, axis=None): return a.sum(axis=axis) summer_wrapped = keepdims_wrapper(summer) assert summer_wrapped != summer assert summer_wrapped == keepdims_wrapper(summer_wrapped) a = np.arange(24).reshape(1, 2, 3, 4) r = summer(a, axis=(1, 3)) rw = summer_wrapped(a, axis=(1, 3), keepdims=True) rwf = summer_wrapped(a, axis=(1, 3), keepdims=False) assert r.ndim == 2 assert r.shape == (1, 3) assert (r == np.array([[60, 92, 124]])).all() assert rw.ndim == 4 assert rw.shape == (1, 1, 3, 1) assert (rw == np.array([[[[60], [92], [124]]]])).all() assert rwf.ndim == 2 assert rwf.shape == (1, 3) assert (rwf == np.array([[60, 92, 124]])).all()
def test_keepdims_wrapper_one_axis(): def summer(a, axis=None): return a.sum(axis=axis) summer_wrapped = keepdims_wrapper(summer) assert summer_wrapped != summer assert summer_wrapped == keepdims_wrapper(summer_wrapped) a = np.arange(24).reshape(1, 2, 3, 4) r = summer(a, axis=2) rw = summer_wrapped(a, axis=2, keepdims=True) rwf = summer_wrapped(a, axis=2, keepdims=False) assert r.ndim == 3 assert r.shape == (1, 2, 4) assert (r == np.array([[[12, 15, 18, 21], [48, 51, 54, 57]]])).all() assert rw.ndim == 4 assert rw.shape == (1, 2, 1, 4) assert (rw == np.array([[[[12, 15, 18, 21]], [[48, 51, 54, 57]]]])).all() assert rwf.ndim == 3 assert rwf.shape == (1, 2, 4) assert (rwf == np.array([[[12, 15, 18, 21], [48, 51, 54, 57]]])).all()