jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def tune():
    with open('convolution.cu', 'r') as f:
        kernel_string = f.read()

    #setup tunable parameters
    tune_params = OrderedDict()
    tune_params["filter_height"] = [i for i in range(3,35,2)]
    tune_params["filter_width"] = [i for i in range(3,35,2)]
    tune_params["block_size_x"] = [16*i for i in range(1,9)]
    tune_params["block_size_y"] = [2**i for i in range(6)]
    tune_params["tile_size_x"] = [i for i in range(1,9)]
    tune_params["tile_size_y"] = [i for i in range(1,9)]

    tune_params["use_padding"] = [0,1]  #toggle the insertion of padding in shared memory
    tune_params["read_only"] = [0,1]    #toggle using the read-only cache

    #limit the search to only use padding when its effective
    restrict = ["use_padding==0 or (block_size_x % 32 != 0)"]

    problem_size = (4096, 4096)
    size = numpy.multiply.reduce(problem_size)
    largest_fh = max(tune_params["filter_height"])
    largest_fw = max(tune_params["filter_width"])
    input_size = ((problem_size[0]+largest_fw-1) * (problem_size[1]+largest_fh-1))

    output_image = numpy.zeros(size).astype(numpy.float32)
    input_image = numpy.random.randn(input_size).astype(numpy.float32)
    filter_weights = numpy.random.randn(largest_fh * largest_fw).astype(numpy.float32)

    cmem_args = {'d_filter': filter_weights}
    args = [output_image, input_image, filter_weights]

    grid_div_x = ["block_size_x", "tile_size_x"]
    grid_div_y = ["block_size_y", "tile_size_y"]

    #start tuning
    import time

    start = time.time()

    results, env = kernel_tuner.tune_kernel("convolution_kernel", kernel_string,
        problem_size, args, tune_params,
        grid_div_y=grid_div_y, grid_div_x=grid_div_x, cmem_args=cmem_args,
        verbose=True, restrictions=restrict)

    end = time.time()
    env['execution_time'] = end-start

    return results, env
//  [python]
def tune():
    with open('convolution.cu', 'r') as f:
        kernel_string = f.read()

    #setup tunable parameters
    tune_params = OrderedDict()
    tune_params["filter_height"] = [i for i in range(3,35,2)]
    tune_params["filter_width"] = [i for i in range(3,35,2)]
    tune_params["block_size_x"] = [16*i for i in range(1,9)]
    tune_params["block_size_y"] = [2**i for i in range(6)]
    tune_params["tile_size_x"] = [i for i in range(1,9)]
    tune_params["tile_size_y"] = [i for i in range(1,9)]

    tune_params["use_padding"] = [0,1]  #toggle the insertion of padding in shared memory
    tune_params["read_only"] = [0,1]    #toggle using the read-only cache

    #limit the search to only use padding when its effective
    restrict = ["use_padding==0 or (block_size_x % 32 != 0)"]

    problem_size = (4096, 4096)
    size = numpy.prod(problem_size)
    largest_fh = max(tune_params["filter_height"])
    largest_fw = max(tune_params["filter_width"])
    input_size = ((problem_size[0]+largest_fw-1) * (problem_size[1]+largest_fh-1))

    output_image = numpy.zeros(size).astype(numpy.float32)
    input_image = numpy.random.randn(input_size).astype(numpy.float32)
    filter_weights = numpy.random.randn(largest_fh * largest_fw).astype(numpy.float32)

    cmem_args = {'d_filter': filter_weights}
    args = [output_image, input_image, filter_weights]

    grid_div_x = ["block_size_x", "tile_size_x"]
    grid_div_y = ["block_size_y", "tile_size_y"]

    #start tuning
    import time

    start = time.time()

    results, env = kernel_tuner.tune_kernel("convolution_kernel", kernel_string,
        problem_size, args, tune_params,
        grid_div_y=grid_div_y, grid_div_x=grid_div_x, cmem_args=cmem_args,
        verbose=True, restrictions=restrict)

    end = time.time()
    env['execution_time'] = end-start

    return results, env
//  [python]
def _get_vpplocG_part1(mydf, with_rho_core=PP_WITH_RHO_CORE):
    cell = mydf.cell
    mesh = mydf.mesh

    if not with_rho_core:
        # compute rho_core directly in G-space
        # this is much slower that the following
        Gv = cell.get_Gv(mesh)
        vpplocG_part1 = get_pp_loc_part1_gs(cell, Gv)
    else:
        # compute rho_core in real space then transform to G-space
        weight = cell.vol / n.prod(mesh)
        rho_core = make_rho_core(cell)
        rhoG_core = weight * tools.fft(rho_core, mesh)
        rho_core = None
        coulG = tools.get_coulG(cell, mesh=mesh)
        vpplocG_part1 = rhoG_core * coulG
        rhoG_core = coulG = None
        # G = 0 contribution
        chargs = cell.atom_charges()
        rloc = []
        for ia in range(cell.natm):
            symb = cell.atom_symbol(ia)
            rloc.append(cell._pseudo[symb][1])
        rloc = numpy.asarray(rloc)
        vpplocG_part1[0] += 2. * numpy.pi * numpy.sum(rloc * rloc * chargs)
    return vpplocG_part1
//  [python]
def _get_vpplocG_part1(mydf, with_rho_core=PP_WITH_RHO_CORE):
    cell = mydf.cell
    mesh = mydf.mesh

    if not with_rho_core:
        # compute rho_core directly in G-space
        # this is much slower that the following
        Gv = cell.get_Gv(mesh)
        vpplocG_part1 = get_pp_loc_part1_gs(cell, Gv)
    else:
        # compute rho_core in real space then transform to G-space
        weight = cell.vol / numpy.prod(mesh)
        rho_core = make_rho_core(cell)
        rhoG_core = weight * tools.fft(rho_core, mesh)
        rho_core = None
        coulG = tools.get_coulG(cell, mesh=mesh)
        vpplocG_part1 = rhoG_core * coulG
        rhoG_core = coulG = None
        # G = 0 contribution
        chargs = cell.atom_charges()
        rloc = []
        for ia in range(cell.natm):
            symb = cell.atom_symbol(ia)
            rloc.append(cell._pseudo[symb][1])
        rloc = numpy.asarray(rloc)
        vpplocG_part1[0] += 2. * numpy.pi * numpy.sum(rloc * rloc * chargs)
    return vpplocG_part1
//  [python]
def make_rho_core(cell, mesh=None, precision=None, atm_id=None):
    if mesh is None:
        mesh = cell.mesh
    fakecell, max_radius = fake_cell_vloc_part1(cell, atm_id=atm_id, precision=precision)
    atm = fakecell._atm
    bas = fakecell._bas
    env = fakecell._env

    a = numpy.asarray(cell.lattice_vectors(), order='C', dtype=float)
    if abs(a - numpy.diag(a.diagonal())).max() < 1e-12:
        lattice_type = '_orth'
    else:
        lattice_type = '_nonorth'
        raise NotImplementedError
    eval_fn = 'make_rho_lda' + lattice_type

    b = numpy.asarray(numpy.linalg.inv(a.T), order='C', dtype=float)
    mesh = numpy.asarray(mesh, order='C', dtype=numpy.int32)
    rho_core = numpy.zeros((n.prod(mesh),), order='C', dtype=float)
    drv = getattr(libdft, 'build_core_density', None)
    try:
        drv(getattr(libdft, eval_fn),
            rho_core.ctypes.data_as(ctypes.c_void_p),
            atm.ctypes.data_as(ctypes.c_void_p),
            bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(len(bas)),
            env.ctypes.data_as(ctypes.c_void_p),
            mesh.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_int(cell.dimension),
            a.ctypes.data_as(ctypes.c_void_p),
            b.ctypes.data_as(ctypes.c_void_p), ctypes.c_double(max_radius))
    except Exception as e:
        raise RuntimeError("Failed to compute rho_core. %s" % e)
    return rho_core
//  [python]
def make_rho_core(cell, mesh=None, precision=None, atm_id=None):
    if mesh is None:
        mesh = cell.mesh
    fakecell, max_radius = fake_cell_vloc_part1(cell, atm_id=atm_id, precision=precision)
    atm = fakecell._atm
    bas = fakecell._bas
    env = fakecell._env

    a = numpy.asarray(cell.lattice_vectors(), order='C', dtype=float)
    if abs(a - numpy.diag(a.diagonal())).max() < 1e-12:
        lattice_type = '_orth'
    else:
        lattice_type = '_nonorth'
        raise NotImplementedError
    eval_fn = 'make_rho_lda' + lattice_type

    b = numpy.asarray(numpy.linalg.inv(a.T), order='C', dtype=float)
    mesh = numpy.asarray(mesh, order='C', dtype=numpy.int32)
    rho_core = numpy.zeros((numpy.prod(mesh),), order='C', dtype=float)
    drv = getattr(libdft, 'build_core_density', None)
    try:
        drv(getattr(libdft, eval_fn),
            rho_core.ctypes.data_as(ctypes.c_void_p),
            atm.ctypes.data_as(ctypes.c_void_p),
            bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(len(bas)),
            env.ctypes.data_as(ctypes.c_void_p),
            mesh.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_int(cell.dimension),
            a.ctypes.data_as(ctypes.c_void_p),
            b.ctypes.data_as(ctypes.c_void_p), ctypes.c_double(max_radius))
    except Exception as e:
        raise RuntimeError("Failed to compute rho_core. %s" % e)
    return rho_core
//  [python]
def reference_layer_norm(inp: np.ndarray, normalized_shape: Tuple[int], weight=None, bias=None, eps=1e-5):
    feature_size = np.multiply.reduce(normalized_shape)
    inp_view = inp.reshape(-1, feature_size)  # type: ignore[call-overload]
    mean = inp_view.mean(axis=-1, keepdims=True)
    var = inp_view.var(axis=-1, ddof=0, keepdims=True)
    Y = (inp_view - mean) / np.sqrt(var + eps)
    if weight is None and bias is not None:
        Y = Y + bias.reshape(-1)
    elif weight is not None and bias is None:
        Y = Y * weight.reshape(-1)
    elif weight is not None and bias is not None:
        Y = Y * weight.reshape(-1) + bias.reshape(-1)
    return Y.reshape(*inp.shape)
//  [python]
def reference_layer_norm(inp: np.ndarray, normalized_shape: Tuple[int], weight=None, bias=None, eps=1e-5):
    feature_size = np.prod(normalized_shape)
    inp_view = inp.reshape(-1, feature_size)  # type: ignore[call-overload]
    mean = inp_view.mean(axis=-1, keepdims=True)
    var = inp_view.var(axis=-1, ddof=0, keepdims=True)
    Y = (inp_view - mean) / np.sqrt(var + eps)
    if weight is None and bias is not None:
        Y = Y + bias.reshape(-1)
    elif weight is not None and bias is None:
        Y = Y * weight.reshape(-1)
    elif weight is not None and bias is not None:
        Y = Y * weight.reshape(-1) + bias.reshape(-1)
    return Y.reshape(*inp.shape)
//  [python]
def generate_inputs(self):
        shape = self.batchsize, self.size
        size = numpy.multiply.reduce(shape) // shape[0]
        x = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
        gamma = numpy.random.uniform(-1, 1, size).astype(self.dtype)
        beta = numpy.random.uniform(-1, 1, size).astype(self.dtype)
        return x, gamma, beta
//  [python]
def generate_inputs(self):
        shape = self.batchsize, self.size
        size = numpy.prod(shape) // shape[0]
        x = numpy.random.uniform(-1, 1, shape).astype(self.dtype)
        gamma = numpy.random.uniform(-1, 1, size).astype(self.dtype)
        beta = numpy.random.uniform(-1, 1, size).astype(self.dtype)
        return x, gamma, beta
//  [python]
def decide_chunk_sizes(shape, chunk_size, itemsize):
    """
    Decide how a given tensor can be split into chunk.

    :param shape: tensor's shape
    :param chunk_size: if dict provided, it's dimension id to chunk size;
                       if provided, it's the chunk size for each dimension.
    :param itemsize: element size
    :return: the calculated chunk size for each dimension
    :rtype: tuple
    """

    from ..config import options

    chunk_size = dictify_chunk_size(shape, chunk_size)
    nleft = len(shape) - len(chunk_size)
    if nleft < 0:
        raise ValueError("chunks have more dimensions than input tensor")
    if nleft == 0:
        return normalize_chunk_sizes(
            shape, tuple(chunk_size[j] for j in range(len(shape)))
        )

    max_chunk_size = options.chunk_store_limit

    # normalize the dimension which specified first
    dim_to_normalized = {
        i: normalize_chunk_sizes((shape[i],), (c,))[0] for i, c in chunk_size.items()
    }

    left = {j: [] for j in range(len(shape)) if j not in dim_to_normalized}
    left_unsplit = {j: shape[j] for j in left}
    while True:
        nbytes_occupied = (
            n.prod([max(c) for c in dim_to_normalized.values()]) * itemsize
        )
        dim_size = np.maximum(
            int(np.power(max_chunk_size / nbytes_occupied, 1 / float(len(left)))), 1
        )
        for j, ns in left.copy().items():
            unsplit = left_unsplit[j]
            ns.append(int(np.minimum(unsplit, dim_size)))
            left_unsplit[j] -= ns[-1]
            if left_unsplit[j] <= 0:
                dim_to_normalized[j] = tuple(ns)
                del left[j]

        if len(left) == 0:
            break

    return tuple(dim_to_normalized[i] for i in range(len(dim_to_normalized)))
//  [python]
def decide_chunk_sizes(shape, chunk_size, itemsize):
    """
    Decide how a given tensor can be split into chunk.

    :param shape: tensor's shape
    :param chunk_size: if dict provided, it's dimension id to chunk size;
                       if provided, it's the chunk size for each dimension.
    :param itemsize: element size
    :return: the calculated chunk size for each dimension
    :rtype: tuple
    """

    from ..config import options

    chunk_size = dictify_chunk_size(shape, chunk_size)
    nleft = len(shape) - len(chunk_size)
    if nleft < 0:
        raise ValueError("chunks have more dimensions than input tensor")
    if nleft == 0:
        return normalize_chunk_sizes(
            shape, tuple(chunk_size[j] for j in range(len(shape)))
        )

    max_chunk_size = options.chunk_store_limit

    # normalize the dimension which specified first
    dim_to_normalized = {
        i: normalize_chunk_sizes((shape[i],), (c,))[0] for i, c in chunk_size.items()
    }

    left = {j: [] for j in range(len(shape)) if j not in dim_to_normalized}
    left_unsplit = {j: shape[j] for j in left}
    while True:
        nbytes_occupied = (
            np.prod([max(c) for c in dim_to_normalized.values()]) * itemsize
        )
        dim_size = np.maximum(
            int(np.power(max_chunk_size / nbytes_occupied, 1 / float(len(left)))), 1
        )
        for j, ns in left.copy().items():
            unsplit = left_unsplit[j]
            ns.append(int(np.minimum(unsplit, dim_size)))
            left_unsplit[j] -= ns[-1]
            if left_unsplit[j] <= 0:
                dim_to_normalized[j] = tuple(ns)
                del left[j]

        if len(left) == 0:
            break

    return tuple(dim_to_normalized[i] for i in range(len(dim_to_normalized)))
//  [python]
def vpploc_part1_nuc_grad(mydf, dm, kpts=numpy.zeros((1,3)), atm_id=None, precision=None):
    from .multigrid_pair import _eval_rhoG
    t0 = (logger.process_clock(), logger.perf_counter())
    cell = mydf.cell
    fakecell, max_radius = fake_cell_vloc_part1(cell, atm_id=atm_id, precision=precision)
    atm = fakecell._atm
    bas = fakecell._bas
    env = fakecell._env

    a = numpy.asarray(cell.lattice_vectors(), order='C', dtype=float)
    if abs(a - numpy.diag(a.diagonal())).max() < 1e-12:
        lattice_type = '_orth'
    else:
        lattice_type = '_nonorth'
        raise NotImplementedError
    eval_fn = 'eval_mat_lda' + lattice_type + '_ip1'

    b = numpy.asarray(numpy.linalg.inv(a.T), order='C', dtype=float)
    mesh = numpy.asarray(mydf.mesh, order='C', dtype=numpy.int32)
    ngrids = numpy.multiply.reduce(mesh)
    comp = 3
    grad = numpy.zeros((len(atm),comp), order="C", dtype=float)
    drv = getattr(libdft, 'int_gauss_charge_v_rs', None)

    if mydf.rhoG is None:
        rhoG = _eval_rhoG(mydf, dm, hermi=1, kpts=kpts, deriv=0)
    else:
        rhoG = mydf.rhoG
    rhoG = rhoG[...,0,:]
    rhoG = rhoG.reshape(-1,ngrids)
    if rhoG.shape[0] == 2: #unrestricted
        rhoG = rhoG[0] + rhoG[1]
    else:
        assert rhoG.shape[0] == 1
        rhoG = rhoG[0]

    coulG = tools.get_coulG(cell, mesh=mesh)
    vG = numpy.multiply(rhoG, coulG)

    v_rs = numpy.asarray(tools.ifft(vG, mesh).real, order="C")
    try:
        drv(getattr(libdft, eval_fn),
            grad.ctypes.data_as(ctypes.c_void_p),
            v_rs.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_int(comp),
            atm.ctypes.data_as(ctypes.c_void_p),
            bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(len(bas)),
            env.ctypes.data_as(ctypes.c_void_p),
            mesh.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_int(cell.dimension),
            a.ctypes.data_as(ctypes.c_void_p),
            b.ctypes.data_as(ctypes.c_void_p), ctypes.c_double(max_radius))
    except Exception as e:
        raise RuntimeError("Failed to computed nuclear gradients of vpploc part1. %s" % e)
    grad *= -1
    t0 = logger.timer(mydf, 'vpploc_part1_nuc_grad', *t0)
    return grad
//  [python]
def vpploc_part1_nuc_grad(mydf, dm, kpts=numpy.zeros((1,3)), atm_id=None, precision=None):
    from .multigrid_pair import _eval_rhoG
    t0 = (logger.process_clock(), logger.perf_counter())
    cell = mydf.cell
    fakecell, max_radius = fake_cell_vloc_part1(cell, atm_id=atm_id, precision=precision)
    atm = fakecell._atm
    bas = fakecell._bas
    env = fakecell._env

    a = numpy.asarray(cell.lattice_vectors(), order='C', dtype=float)
    if abs(a - numpy.diag(a.diagonal())).max() < 1e-12:
        lattice_type = '_orth'
    else:
        lattice_type = '_nonorth'
        raise NotImplementedError
    eval_fn = 'eval_mat_lda' + lattice_type + '_ip1'

    b = numpy.asarray(numpy.linalg.inv(a.T), order='C', dtype=float)
    mesh = numpy.asarray(mydf.mesh, order='C', dtype=numpy.int32)
    ngrids = numpy.prod(mesh)
    comp = 3
    grad = numpy.zeros((len(atm),comp), order="C", dtype=float)
    drv = getattr(libdft, 'int_gauss_charge_v_rs', None)

    if mydf.rhoG is None:
        rhoG = _eval_rhoG(mydf, dm, hermi=1, kpts=kpts, deriv=0)
    else:
        rhoG = mydf.rhoG
    rhoG = rhoG[...,0,:]
    rhoG = rhoG.reshape(-1,ngrids)
    if rhoG.shape[0] == 2: #unrestricted
        rhoG = rhoG[0] + rhoG[1]
    else:
        assert rhoG.shape[0] == 1
        rhoG = rhoG[0]

    coulG = tools.get_coulG(cell, mesh=mesh)
    vG = numpy.multiply(rhoG, coulG)

    v_rs = numpy.asarray(tools.ifft(vG, mesh).real, order="C")
    try:
        drv(getattr(libdft, eval_fn),
            grad.ctypes.data_as(ctypes.c_void_p),
            v_rs.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_int(comp),
            atm.ctypes.data_as(ctypes.c_void_p),
            bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(len(bas)),
            env.ctypes.data_as(ctypes.c_void_p),
            mesh.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_int(cell.dimension),
            a.ctypes.data_as(ctypes.c_void_p),
            b.ctypes.data_as(ctypes.c_void_p), ctypes.c_double(max_radius))
    except Exception as e:
        raise RuntimeError("Failed to computed nuclear gradients of vpploc part1. %s" % e)
    grad *= -1
    t0 = logger.timer(mydf, 'vpploc_part1_nuc_grad', *t0)
    return grad
//  [python]
def test_basic(self):
        ba = [1, 2, 10, 11, 6, 5, 4]
        ba2 = [[1, 2, 3, 4], [5, 6, 7, 9], [10, 3, 4, 5]]
        for ctype in [np.int16, np.uint16, np.int32, np.uint32,
                      np.float32, np.float64, np.complex64, np.complex128]:
            a = np.array(ba, ctype)
            a2 = np.array(ba2, ctype)
            if ctype in ['1', 'b']:
                self.assertRaises(ArithmeticError, np.prod, a)
                self.assertRaises(ArithmeticError, np.prod, a2, 1)
                self.assertRaises(ArithmeticError, np.prod, a)
            else:
                assert_equal(np.prod(a, axis=0), 26400)
                assert_array_equal(np.prod(a2, axis=0),
                                   np.array([50, 36, 84, 180], ctype))
                assert_array_equal(np.prod(a2, axis=-1),
                                   np.array([24, 1890, 600], ctype))
//  [python]
def normal(self, size, avg=0.0, std=1.0, ndim=None,
               dtype=None, nstreams=None):
        # TODO : need description for method
        """
        Parameters
        ----------
        size
            Can be a list of integers or Theano variables (ex: the shape
            of another Theano Variable).
        dtype
            The output data type. If dtype is not specified, it will be
            inferred from the dtype of low and high, but will be at
            least as precise as floatX.
        nstreams
            Number of streams.

        """
        # We need an even number of ]0,1[ samples. Then we split them
        # in two halves. First half becomes our U1's for Box-Muller,
        # second half our U2's. See Wikipedia page:
        # http://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform
        avg = as_tensor_variable(avg)
        std = as_tensor_variable(std)

        if dtype is None:
            dtype = scal.upcast(config.floatX, avg.dtype, std.dtype)

        avg = cast(avg, dtype)
        std = cast(std, dtype)

        evened = False
        constant = False
        if (isinstance(size, tuple) and
                all([isinstance(i, (numpy.integer, integer_types)) for i in size])):
            constant = True
            # Force dtype because it defaults to float when size is empty
            n_samples = numpy.prod(size, dtype='int64')

            if n_samples % 2 == 1:
                n_samples += 1
                evened = True
        else:
            # if even, don't change, if odd, +1
            n_samples = n.prod(size) + (n.prod(size) % 2)
        flattened = self.uniform(size=(n_samples,), dtype=dtype,
                                 nstreams=nstreams)

        if constant:
            U1 = flattened[:n_samples // 2]
            U2 = flattened[n_samples // 2:]
        else:
            U1 = flattened[:n.prod(flattened.shape) // 2]
            U2 = flattened[n.prod(flattened.shape) // 2:]

        # normal_samples = zeros_like(flattened)
        sqrt_ln_U1 = sqrt(-2.0 * log(U1))
        # TypeError: 'TensorVariable' object does not support item assignment
        # so this doesn't work...
        # normal_samples[:n_samples/2] = sqrt_ln_U1 * cos(2.0*numpy.pi*U2)
        # normal_samples[n_samples/2:] = sqrt_ln_U1 * sin(2.0*numpy.pi*U2)

        # so trying this instead
        first_half = sqrt_ln_U1 * cos(
            numpy.array(2.0 * numpy.pi, dtype=dtype) * U2)
        second_half = sqrt_ln_U1 * sin(
            numpy.array(2.0 * numpy.pi, dtype=dtype) * U2)
        normal_samples = join(0, first_half, second_half)

        final_samples = None
        if evened:
            final_samples = normal_samples[:-1]
        elif constant:
            final_samples = normal_samples
        else:
            final_samples = normal_samples[:n.prod(size)]

        if not size:
            # Force the dtype to be int64, otherwise reshape complains
            size = tensor.constant(size, dtype='int64')
        final_samples = final_samples.reshape(size)

        final_samples = avg + std * final_samples

        assert final_samples.dtype == dtype
        return final_samples
//  [python]
def normal(self, size, avg=0.0, std=1.0, ndim=None,
               dtype=None, nstreams=None):
        # TODO : need description for method
        """
        Parameters
        ----------
        size
            Can be a list of integers or Theano variables (ex: the shape
            of another Theano Variable).
        dtype
            The output data type. If dtype is not specified, it will be
            inferred from the dtype of low and high, but will be at
            least as precise as floatX.
        nstreams
            Number of streams.

        """
        # We need an even number of ]0,1[ samples. Then we split them
        # in two halves. First half becomes our U1's for Box-Muller,
        # second half our U2's. See Wikipedia page:
        # http://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform
        avg = as_tensor_variable(avg)
        std = as_tensor_variable(std)

        if dtype is None:
            dtype = scal.upcast(config.floatX, avg.dtype, std.dtype)

        avg = cast(avg, dtype)
        std = cast(std, dtype)

        evened = False
        constant = False
        if (isinstance(size, tuple) and
                all([isinstance(i, (numpy.integer, integer_types)) for i in size])):
            constant = True
            # Force dtype because it defaults to float when size is empty
            n_samples = numpy.prod(size, dtype='int64')

            if n_samples % 2 == 1:
                n_samples += 1
                evened = True
        else:
            # if even, don't change, if odd, +1
            n_samples = prod(size) + (prod(size) % 2)
        flattened = self.uniform(size=(n_samples,), dtype=dtype,
                                 nstreams=nstreams)

        if constant:
            U1 = flattened[:n_samples // 2]
            U2 = flattened[n_samples // 2:]
        else:
            U1 = flattened[:prod(flattened.shape) // 2]
            U2 = flattened[prod(flattened.shape) // 2:]

        # normal_samples = zeros_like(flattened)
        sqrt_ln_U1 = sqrt(-2.0 * log(U1))
        # TypeError: 'TensorVariable' object does not support item assignment
        # so this doesn't work...
        # normal_samples[:n_samples/2] = sqrt_ln_U1 * cos(2.0*numpy.pi*U2)
        # normal_samples[n_samples/2:] = sqrt_ln_U1 * sin(2.0*numpy.pi*U2)

        # so trying this instead
        first_half = sqrt_ln_U1 * cos(
            numpy.array(2.0 * numpy.pi, dtype=dtype) * U2)
        second_half = sqrt_ln_U1 * sin(
            numpy.array(2.0 * numpy.pi, dtype=dtype) * U2)
        normal_samples = join(0, first_half, second_half)

        final_samples = None
        if evened:
            final_samples = normal_samples[:-1]
        elif constant:
            final_samples = normal_samples
        else:
            final_samples = normal_samples[:prod(size)]

        if not size:
            # Force the dtype to be int64, otherwise reshape complains
            size = tensor.constant(size, dtype='int64')
        final_samples = final_samples.reshape(size)

        final_samples = avg + std * final_samples

        assert final_samples.dtype == dtype
        return final_samples
//  [python]
def setUpClass(cls):
        # Choose names
        cls.dvid_repo = "datasetA"
        cls.data_name = "random_data"
        cls.volume_location = "/repos/{dvid_repo}/volumes/{data_name}".format(**cls.__dict__)

        cls.data_uuid = get_testrepo_root_uuid()
        cls.node_location = "/repos/{dvid_repo}/nodes/{data_uuid}".format(**cls.__dict__)

        # Generate some test data
        # data = numpy.random.randint(0, 255, (128, 256, 512, 1)).astype( numpy.uint8 )

        data = numpy.zeros((128, 256, 512, 1), dtype=numpy.uint8)
        data.flat[:] = numpy.arange(n.prod((128, 256, 512, 1)))
        cls.original_data = data
        cls.voxels_metadata = VoxelsMetadata.create_default_metadata(data.shape, data.dtype, "zyxc", 1.0, "")

        # Write it to a new data instance
        node_service = DVIDNodeService(TEST_DVID_SERVER, cls.data_uuid)

        node_service.create_grayscale8(cls.data_name)
        node_service.put_gray3D(cls.data_name, data[..., 0], (0, 0, 0))
//  [python]
def setUpClass(cls):
        # Choose names
        cls.dvid_repo = "datasetA"
        cls.data_name = "random_data"
        cls.volume_location = "/repos/{dvid_repo}/volumes/{data_name}".format(**cls.__dict__)

        cls.data_uuid = get_testrepo_root_uuid()
        cls.node_location = "/repos/{dvid_repo}/nodes/{data_uuid}".format(**cls.__dict__)

        # Generate some test data
        # data = numpy.random.randint(0, 255, (128, 256, 512, 1)).astype( numpy.uint8 )

        data = numpy.zeros((128, 256, 512, 1), dtype=numpy.uint8)
        data.flat[:] = numpy.arange(numpy.prod((128, 256, 512, 1)))
        cls.original_data = data
        cls.voxels_metadata = VoxelsMetadata.create_default_metadata(data.shape, data.dtype, "zyxc", 1.0, "")

        # Write it to a new data instance
        node_service = DVIDNodeService(TEST_DVID_SERVER, cls.data_uuid)

        node_service.create_grayscale8(cls.data_name)
        node_service.put_gray3D(cls.data_name, data[..., 0], (0, 0, 0))
//  [python]
def reduce_prod(
        cls, tensor: XTensor, axis: Optional[List[int]] = None, keepdims: bool = False
    ) -> XTensor:
        assert tensor is not None, "'tensor' should not be None."
        assert isinstance(tensor, XTensor), "'tensor' should be an XTensor instance."

        in_tensor = tensor.to_numpy()
        out_tensor = np.prod(
            in_tensor, axis=tuple(axis), keepdims=keepdims, dtype=in_tensor.dtype
        )

        if not isinstance(out_tensor, np.ndarray):
            if out_tensor.__class__ in [np.int32, np.int64]:
                out_tensor = np.array([out_tensor])
            else:
                raise NotImplementedError(
                    f"[ERROR] Not supported data type: {out_tensor.__class__.__name__}"
                )

        return XTensor(out_tensor, format=tensor.data_format)
//  [python]
def test_fft_kpq(nalpha):
    # Create regular grid.
    nmax = 2
    mesh = [2*nmax+1]*2
    grid, eigs = generate_fft_grid(mesh)
    
    qmax = 2*nmax
    qmesh = [2*qmax+1]*2
    qgrid, qeigs = generate_fft_grid(qmesh)

    # Create wavefunction
    nbasis = len(grid)
    
    numpy.random.seed(7)

    psi = get_random_wavefunction((nalpha,nalpha), nbasis)
    I = get_random_wavefunction((nalpha,nalpha), nbasis)
    trial = I[:,:nalpha].conj()
    
    G, Gh = gab_mod(trial,psi[:,:nalpha])
    nqgrid = n.prod(qmesh)

    # Check by direct convolution f(q) = \sum_G Psi[G+Q] Gh[G].
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for i, q in enumerate(qgrid):
        for j, k in enumerate(grid):
            kpq = k+q
            ikpq = lookup(kpq, grid)
            if ikpq is not None:
                fq_direct[i] += trial[ikpq,0] * Gh[0,j]

    trial_grid = numpy.flip(trial[:,0]).reshape(mesh)
    Gh_grid = Gh[0,:].reshape(mesh)

    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    fq_conv += nqgrid*convolve(Gh_grid, trial_grid,mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal
    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(Gh_grid,trial_grid).ravel())

    import matplotlib.pyplot as pl
    pl.plot(fq_conv, label='fft')
    pl.plot(fq_conv_sc, label='fft_scipy')
    pl.plot(fq_direct, label='direct')
    # pl.plot(fq, label='from_gf')
    pl.legend()
    pl.show()
//  [python]
def test_fft_kpq(nalpha):
    # Create regular grid.
    nmax = 2
    mesh = [2*nmax+1]*2
    grid, eigs = generate_fft_grid(mesh)
    
    qmax = 2*nmax
    qmesh = [2*qmax+1]*2
    qgrid, qeigs = generate_fft_grid(qmesh)

    # Create wavefunction
    nbasis = len(grid)
    
    numpy.random.seed(7)

    psi = get_random_wavefunction((nalpha,nalpha), nbasis)
    I = get_random_wavefunction((nalpha,nalpha), nbasis)
    trial = I[:,:nalpha].conj()
    
    G, Gh = gab_mod(trial,psi[:,:nalpha])
    nqgrid = numpy.prod(qmesh)

    # Check by direct convolution f(q) = \sum_G Psi[G+Q] Gh[G].
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for i, q in enumerate(qgrid):
        for j, k in enumerate(grid):
            kpq = k+q
            ikpq = lookup(kpq, grid)
            if ikpq is not None:
                fq_direct[i] += trial[ikpq,0] * Gh[0,j]

    trial_grid = numpy.flip(trial[:,0]).reshape(mesh)
    Gh_grid = Gh[0,:].reshape(mesh)

    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    fq_conv += nqgrid*convolve(Gh_grid, trial_grid,mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal
    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(Gh_grid,trial_grid).ravel())

    import matplotlib.pyplot as pl
    pl.plot(fq_conv, label='fft')
    pl.plot(fq_conv_sc, label='fft_scipy')
    pl.plot(fq_direct, label='direct')
    # pl.plot(fq, label='from_gf')
    pl.legend()
    pl.show()
//  [python]
def tune():
    with open("convolution.cu", "r") as f:
        kernel_string = f.read()

    filter_size = (17, 17)
    problem_size = (4096, 4096)
    size = numpy.prod(problem_size)
    border_size = (filter_size[0] // 2 * 2, filter_size[1] // 2 * 2)
    input_size = (problem_size[0] + border_size[0]) * (problem_size[1] + border_size[1])

    output = numpy.zeros(size).astype(numpy.float32)
    input = numpy.random.randn(input_size).astype(numpy.float32)

    filter = numpy.random.randn(filter_size[0] * filter_size[1]).astype(numpy.float32)
    cmem_args = {"d_filter": filter}

    args = [output, input, filter]
    tune_params = OrderedDict()
    tune_params["filter_width"] = [filter_size[0]]
    tune_params["filter_height"] = [filter_size[1]]

    # tune_params["block_size_x"] = [16*i for i in range(1,3)]
    tune_params["block_size_x"] = [16 * i for i in range(1, 9)]
    # tune_params["block_size_y"] = [2**i for i in range(1,5)]
    tune_params["block_size_y"] = [2**i for i in range(1, 6)]

    tune_params["tile_size_x"] = [2**i for i in range(3)]
    tune_params["tile_size_y"] = [2**i for i in range(3)]

    tune_params["use_padding"] = [
        0,
        1,
    ]  # toggle the insertion of padding in shared memory
    tune_params["read_only"] = [0, 1]  # toggle using the read-only cache

    grid_div_x = ["block_size_x", "tile_size_x"]
    grid_div_y = ["block_size_y", "tile_size_y"]

    # compute the answer using a naive kernel
    params = {"block_size_x": 16, "block_size_y": 16}
    tune_params["filter_width"] = [filter_size[0]]
    tune_params["filter_height"] = [filter_size[1]]
    results = kernel_tuner.run_kernel(
        "convolution_naive",
        kernel_string,
        problem_size,
        args,
        params,
        grid_div_y=["block_size_y"],
        grid_div_x=["block_size_x"],
        lang="cupy",
    )

    # set non-output fields to None
    answer = [results[0], None, None]

    # start kernel tuning with correctness verification
    return kernel_tuner.tune_kernel(
        "convolution_kernel",
        kernel_string,
        problem_size,
        args,
        tune_params,
        grid_div_y=grid_div_y,
        grid_div_x=grid_div_x,
        verbose=True,
        cmem_args=cmem_args,
        answer=answer,
        lang="cupy",
    )
//  [python]
def open_mmap(self, shape=None, mode='r', dtype='float32', **kwargs):
        '''
        Opens a numpy memmap at filename [blob][-part].mmap.
        Defaults to read-only (override with mode='w+').
        Defaults to float32 data (override with dtype='int32').
        Defaults to linear shape (override with shape=(...)).
        If the supplied shape contains a -1, then infers the missing
        dimension from the size of the file.
        '''
        knownshape = None
        if shape is not None and not isinstance(shape, (list, tuple)):
            shape = (shape,)
        if shape is not None and -1 not in shape:
            knownshape = shape
        nameargs = dict(kwargs)
        if 'inc' not in nameargs:
            nameargs['inc'] = (mode and 'w' in mode)
        aspair = ('aspair' in nameargs) and nameargs['aspair']
        if aspair:
            num, fn = self.mmap_filename(**nameargs)
        else:
            fn = self.mmap_filename(**nameargs)
        result = numpy.memmap(
                fn,
                mode=mode,
                dtype=dtype,
                shape=knownshape)
        # Supports inferring shape of one unknown dimension
        if shape is not None and -1 in shape:
            lastdim = len(result) // abs(numpy.prod(shape))
            inferred = tuple(map(lambda x: x if x >= 0 else lastdim, shape))
            result.shape = inferred
        if aspair:
            return num, result
        else:
            return result
//  [python]
def iou(candidates: np.ndarray, objects: np.ndarray) -> np.ndarray:
    """
    Calculate IoU between two sets of bounding boxes. Both sets of boxes are expected
    to be in `[x_min, y_min, x_max, y_max]` format.

    Normal IoU is 1 when the boxes are the same and 0 when they don't overlap,
    to transform that into a distance that makes sense we return `1 - iou`.

    Parameters
    ----------
    candidates : numpy.ndarray
        (N, 4) numpy.ndarray containing candidates bounding boxes.
    objects : numpy.ndarray
        (K, 4) numpy.ndarray containing objects bounding boxes.

    Returns
    -------
    numpy.ndarray
        (N, K) numpy.ndarray of `1 - iou` between candidates and objects.
    """
    _validate_bboxes(candidates)

    area_candidates = _boxes_area(candidates.T)
    area_objects = _boxes_area(objects.T)

    top_left = np.maximum(candidates[:, None, :2], objects[:, :2])
    bottom_right = np.minimum(candidates[:, None, 2:], objects[:, 2:])

    area_intersection = np.multiply.reduce(
        np.clip(bottom_right - top_left, a_min=0, a_max=None), 2
    )
    return 1 - area_intersection / (
        area_candidates[:, None] + area_objects - area_intersection
    )
//  [python]
def iou(candidates: np.ndarray, objects: np.ndarray) -> np.ndarray:
    """
    Calculate IoU between two sets of bounding boxes. Both sets of boxes are expected
    to be in `[x_min, y_min, x_max, y_max]` format.

    Normal IoU is 1 when the boxes are the same and 0 when they don't overlap,
    to transform that into a distance that makes sense we return `1 - iou`.

    Parameters
    ----------
    candidates : numpy.ndarray
        (N, 4) numpy.ndarray containing candidates bounding boxes.
    objects : numpy.ndarray
        (K, 4) numpy.ndarray containing objects bounding boxes.

    Returns
    -------
    numpy.ndarray
        (N, K) numpy.ndarray of `1 - iou` between candidates and objects.
    """
    _validate_bboxes(candidates)

    area_candidates = _boxes_area(candidates.T)
    area_objects = _boxes_area(objects.T)

    top_left = np.maximum(candidates[:, None, :2], objects[:, :2])
    bottom_right = np.minimum(candidates[:, None, 2:], objects[:, 2:])

    area_intersection = np.prod(
        np.clip(bottom_right - top_left, a_min=0, a_max=None), 2
    )
    return 1 - area_intersection / (
        area_candidates[:, None] + area_objects - area_intersection
    )
//  [python]
def interpolate_G(self, real_batch, step=0, root_path='.', train_epoch=0):
        batch_size = len(real_batch)
        half_batch_size = int(batch_size/2)

        self.sess.run(self.z_r_update)
        tf_real_batch = to_nchw_numpy(real_batch)
        for i in trange(train_epoch):
            z_r_loss, _ = self.sess.run([self.z_r_loss, self.z_r_optim], {self.x: tf_real_batch})
        z = self.sess.run(self.z_r)

        z1, z2 = z[:half_batch_size], z[half_batch_size:]
        real1_batch, real2_batch = real_batch[:half_batch_size], real_batch[half_batch_size:]

        generated = []
        for idx, ratio in enumerate(np.linspace(0, 1, 10)):
            z = np.stack([slerp(ratio, r1, r2) for r1, r2 in zip(z1, z2)])
            z_decode = self.generate(z, save=False)
            generated.append(z_decode)

        generated = np.stack(generated).transpose([1, 0, 2, 3, 4])
        for idx, img in enumerate(generated):
            save_image(img, os.path.join(root_path, 'test{}_interp_G_{}.png'.format(step, idx)), nrow=10)

        all_img_num = np.multiply.reduce(generated.shape[:2])
        batch_generated = np.reshape(generated, [all_img_num] + list(generated.shape[2:]))
        save_image(batch_generated, os.path.join(root_path, 'test{}_interp_G.png'.format(step)), nrow=10)
//  [python]
def interpolate_G(self, real_batch, step=0, root_path='.', train_epoch=0):
        batch_size = len(real_batch)
        half_batch_size = int(batch_size/2)

        self.sess.run(self.z_r_update)
        tf_real_batch = to_nchw_numpy(real_batch)
        for i in trange(train_epoch):
            z_r_loss, _ = self.sess.run([self.z_r_loss, self.z_r_optim], {self.x: tf_real_batch})
        z = self.sess.run(self.z_r)

        z1, z2 = z[:half_batch_size], z[half_batch_size:]
        real1_batch, real2_batch = real_batch[:half_batch_size], real_batch[half_batch_size:]

        generated = []
        for idx, ratio in enumerate(np.linspace(0, 1, 10)):
            z = np.stack([slerp(ratio, r1, r2) for r1, r2 in zip(z1, z2)])
            z_decode = self.generate(z, save=False)
            generated.append(z_decode)

        generated = np.stack(generated).transpose([1, 0, 2, 3, 4])
        for idx, img in enumerate(generated):
            save_image(img, os.path.join(root_path, 'test{}_interp_G_{}.png'.format(step, idx)), nrow=10)

        all_img_num = np.prod(generated.shape[:2])
        batch_generated = np.reshape(generated, [all_img_num] + list(generated.shape[2:]))
        save_image(batch_generated, os.path.join(root_path, 'test{}_interp_G.png'.format(step)), nrow=10)
//  [python]
def test_fft_kmq(nalpha):
    # Create regular grid.
    nmax = 2
    mesh = [2*nmax+1]*2
    grid, eigs = generate_fft_grid(mesh)
    
    qmax = 2*nmax
    qmesh = [2*qmax+1]*2
    qgrid, qeigs = generate_fft_grid(qmesh)

    # Create wavefunction
    nbasis = len(grid)
    
    numpy.random.seed(7)

    psi = get_random_wavefunction((nalpha,nalpha), nbasis)
    I = get_random_wavefunction((nalpha,nalpha), nbasis)
    trial = I[:,:nalpha].conj()
    
    G, Gh = gab_mod(trial,psi[:,:nalpha])
    nqgrid = numpy.multiply.reduce(qmesh)

    # Check by direct convolution f(q) = \sum_G Psi[G+Q] Gh[G].
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for i, q in enumerate(qgrid):
        for j, k in enumerate(grid):
            kmq = k-q
            ikmq = lookup(kmq, grid)
            if ikmq is not None:
                fq_direct[i] += trial[ikmq,0] * Gh[0,j]

    trial_grid = trial[:,0].reshape(mesh)
    Gh_grid = numpy.flip(Gh[0,:]).reshape(mesh)

    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    fq_conv += nqgrid*convolve(trial_grid, Gh_grid, mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal
    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(trial_grid, Gh_grid).ravel())

    import matplotlib.pyplot as pl
    pl.plot(fq_conv, label='fft')
    pl.plot(fq_conv_sc, label='fft_scipy')
    pl.plot(fq_direct, label='direct')
    # pl.plot(fq, label='from_gf')
    pl.legend()
    pl.show()
//  [python]
def test_fft_kmq(nalpha):
    # Create regular grid.
    nmax = 2
    mesh = [2*nmax+1]*2
    grid, eigs = generate_fft_grid(mesh)
    
    qmax = 2*nmax
    qmesh = [2*qmax+1]*2
    qgrid, qeigs = generate_fft_grid(qmesh)

    # Create wavefunction
    nbasis = len(grid)
    
    numpy.random.seed(7)

    psi = get_random_wavefunction((nalpha,nalpha), nbasis)
    I = get_random_wavefunction((nalpha,nalpha), nbasis)
    trial = I[:,:nalpha].conj()
    
    G, Gh = gab_mod(trial,psi[:,:nalpha])
    nqgrid = numpy.prod(qmesh)

    # Check by direct convolution f(q) = \sum_G Psi[G+Q] Gh[G].
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for i, q in enumerate(qgrid):
        for j, k in enumerate(grid):
            kmq = k-q
            ikmq = lookup(kmq, grid)
            if ikmq is not None:
                fq_direct[i] += trial[ikmq,0] * Gh[0,j]

    trial_grid = trial[:,0].reshape(mesh)
    Gh_grid = numpy.flip(Gh[0,:]).reshape(mesh)

    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    fq_conv += nqgrid*convolve(trial_grid, Gh_grid, mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal
    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(trial_grid, Gh_grid).ravel())

    import matplotlib.pyplot as pl
    pl.plot(fq_conv, label='fft')
    pl.plot(fq_conv_sc, label='fft_scipy')
    pl.plot(fq_direct, label='direct')
    # pl.plot(fq, label='from_gf')
    pl.legend()
    pl.show()
//  [python]
def get_total_subsampling_factor(self):
        """Get total subsampling factor."""
        if self.shared_encoder:
            return self.shared_encoder.conv_subsampling_factor * int(
                n.prod(self.subsample)
        )
        else:
            return self.chn_encoder.conv_subsampling_factor * int(
                n.prod(self.subsample)
        )
//  [python]
def get_total_subsampling_factor(self):
        """Get total subsampling factor."""
        if self.shared_encoder:
            return self.shared_encoder.conv_subsampling_factor * int(
                numpy.prod(self.subsample)
        )
        else:
            return self.chn_encoder.conv_subsampling_factor * int(
                numpy.prod(self.subsample)
        )
//  [python]
def size(self):
        """number of elements in the NDFrame"""
        return np.prod(self.shape)
//  [python]
def __invert__(self):
        try:
            arr = operator.inv(_values_from_object(self))
            return self.__array_wrap__(arr)
        except:

            # inv fails with 0 len
            if not numpy.prod(self.shape):
                return self

            raise
//  [python]
def __invert__(self):
        try:
            arr = operator.inv(_values_from_object(self))
            return self.__array_wrap__(arr)
        except:

            # inv fails with 0 len
            if not np.prod(self.shape):
                return self

            raise
//  [python]
def _pre_ext_drift(self, pnt_cnt, ext_drift=None, set_cond=False):
        """
        Preprocessor for external drifts.

        Parameters
        ----------
        pnt_cnt : :class:`numpy.ndarray`
            Number of points of the mesh.
        ext_drift : :class:`numpy.ndarray` or :any:`None`, optional
            the external drift values at the given positions (only for EDK)
            For multiple external drifts, the first dimension
            should be the index of the drift term.
        set_cond : :class:`bool`, optional
            State if the given external drift is set for the conditioning
            points. Default: False

        Returns
        -------
        ext_drift : :class:`numpy.ndarray` or :any:`None`
            the drift values at the given positions
        """
        if ext_drift is not None:
            ext_drift = np.atleast_2d(np.asarray(ext_drift, dtype=np.double))
            if ext_drift.size == 0:  # treat empty array as no ext_drift
                return np.array([])
            if set_cond:
                if len(ext_drift.shape) > 2 or ext_drift.shape[1] != pnt_cnt:
                    raise ValueError("Krige: wrong number of ext. drifts.")
                return ext_drift
            ext_shape = np.shape(ext_drift)
            shape = (self.ext_drift_no, pnt_cnt)
            if self.drift_no > 1 and ext_shape[0] != self.ext_drift_no:
                raise ValueError("Krige: wrong number of external drifts.")
            if np.prod(ext_shape) != np.prod(shape):
                raise ValueError("Krige: wrong number of ext. drift values.")
            return np.asarray(ext_drift, dtype=np.double).reshape(shape)
        if not set_cond and self._cond_ext_drift.size > 0:
            raise ValueError("Krige: wrong number of ext. drift values.")
        return np.array([])
//  [python]
def size(self) -> int:
        return numpy.prod(self.shape)
//  [python]
def test_enumerating_directions():
    for backend in imp_op_backends:
        print("testing directions for", backend.framework_name)
        for shape in [[], [1], [1, 1, 1], [2, 3, 5, 7]]:
            x = numpy.arange(numpy.prod(shape)).reshape(shape)
            axes1 = _enumerate_directions(x)
            axes2 = _enumerate_directions(backend.from_numpy(x))
            assert len(axes1) == len(axes2) == len(shape)
            for ax1, ax2 in zip(axes1, axes2):
                ax2 = backend.to_numpy(ax2)
                assert ax1.shape == ax2.shape
                assert numpy.allclose(ax1, ax2)