cupy_pr2822.py

import itertools

import numpy

import cupy
import cupyx


shapes = [(4096, 4096), (64, 64, 64, 64)]


for shape in shapes:
    ndim = len(shape)
    size = numpy.prod(shape)
    range_ndim = list(range(ndim))
    for perm in itertools.permutations(range_ndim, ndim):
        for n_axis in range(1, ndim + 1):
            for axes in itertools.combinations(range_ndim, n_axis):
                x = cupy.arange(size, dtype='f').reshape(shape)
                x = x.transpose(perm)
                perf = cupyx.time.repeat(cupy.sum, (x, axes), n=100)
                print(perf)