test_ops_reduction.py 文件源码

def test_reduce_mean():
    data = np.array([[[5, 1], [20, 2]], [[30, 1], [40, 2]], [[55, 1], [60, 2]]], dtype=np.float32)

    assert np.array_equal(import_and_compute('ReduceMean', data), np.mean(data, keepdims=True))
    assert np.array_equal(import_and_compute('ReduceMean', data, keepdims=0),
                          np.mean(data, keepdims=False))

    assert np.array_equal(import_and_compute('ReduceMean', data, axes=(1,)),
                          np.mean(data, keepdims=True, axis=(1,)))
    assert np.array_equal(import_and_compute('ReduceMean', data, axes=(1,), keepdims=0),
                          np.mean(data, keepdims=False, axis=(1,)))

    assert np.array_equal(import_and_compute('ReduceMean', data, axes=(0, 2)),
                          np.mean(data, keepdims=True, axis=(0, 2)))
    assert np.array_equal(import_and_compute('ReduceMean', data, axes=(0, 2), keepdims=0),
                          np.mean(data, keepdims=False, axis=(0, 2)))

    assert np.array_equal(import_and_compute('ReduceMean', data, axes=(0, 1, 2)),
                          np.mean(data, keepdims=True, axis=(0, 1, 2)))
    assert np.array_equal(import_and_compute('ReduceMean', data, axes=(0, 1, 2), keepdims=0),
                          np.mean(data, keepdims=False, axis=(0, 1, 2)))