jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _get_indices_dict(label_list, keys):
    shape = list(map(len, keys))

    group_index = get_group_index(label_list, shape, sort=True, xnull=True)
    ngroups = ((group_index.size and group_index.max()) + 1) \
        if _int64_overflow_possible(shape) \
        else np.prod(shape, dtype='i8')

    sorter = _get_group_index_sorter(group_index, ngroups)

    sorted_labels = [lab.take(sorter) for lab in label_list]
    group_index = group_index.take(sorter)

    return lib.indices_fast(sorter, group_index, keys, sorted_labels)
//  [python]
def get_channel_rmse(self, image1, image2):
        """
        Get the channel-by-channel root mean square error.

        Args:
            image1 (str): Path to the first image.
            image2 (str): Path to the second image.

        Returns:
            tuple: Tuple of floats containing the RMSE per channel.
            None: Images could not be compared because they are different size.
        """
        img1 = Image.open(image1)
        img2 = Image.open(image2)

        # Images must have same size
        if img1.size != img2.size:
            return None

        # Images must have same number of color channels
        if img1.getbands() != img2.getbands():
            # ... unless the only different is alpha
            self.assertEqual(img1.getbands(), ("R", "G", "B"))
            self.assertEqual(img2.getbands(), ("R", "G", "B", "A"))

            # ... and the alpha is always one
            bands = img2.split()
            alphaHist = bands[3].histogram()
            self.assertEqual(sum(alphaHist[:-1]), 0)

            # Generate a version of img2 without alpha
            img2 = Image.merge("RGB", (bands[0], bands[1], bands[2]))

        # Compute root mean square error
        img1bands = img1.split()
        img2bands = img2.split()

        rmseVals = []
        imgBands = zip(img1bands, img2bands)
        for img1Ch, img2Ch in imgBands:
            imSz = numpy.multiply.reduce(img1Ch.size)
            dat1 = numpy.array(img1Ch)
            dat2 = numpy.array(img2Ch)

            sad = numpy.sum(numpy.square(dat1 - dat2))
            mse = numpy.divide(sad, imSz)
            rmse = numpy.sqrt(mse)
            rmseVals.append(rmse)

        return rmseVals
//  [python]
def get_channel_rmse(self, image1, image2):
        """
        Get the channel-by-channel root mean square error.

        Args:
            image1 (str): Path to the first image.
            image2 (str): Path to the second image.

        Returns:
            tuple: Tuple of floats containing the RMSE per channel.
            None: Images could not be compared because they are different size.
        """
        img1 = Image.open(image1)
        img2 = Image.open(image2)

        # Images must have same size
        if img1.size != img2.size:
            return None

        # Images must have same number of color channels
        if img1.getbands() != img2.getbands():
            # ... unless the only different is alpha
            self.assertEqual(img1.getbands(), ("R", "G", "B"))
            self.assertEqual(img2.getbands(), ("R", "G", "B", "A"))

            # ... and the alpha is always one
            bands = img2.split()
            alphaHist = bands[3].histogram()
            self.assertEqual(sum(alphaHist[:-1]), 0)

            # Generate a version of img2 without alpha
            img2 = Image.merge("RGB", (bands[0], bands[1], bands[2]))

        # Compute root mean square error
        img1bands = img1.split()
        img2bands = img2.split()

        rmseVals = []
        imgBands = zip(img1bands, img2bands)
        for img1Ch, img2Ch in imgBands:
            imSz = numpy.prod(img1Ch.size)
            dat1 = numpy.array(img1Ch)
            dat2 = numpy.array(img2Ch)

            sad = numpy.sum(numpy.square(dat1 - dat2))
            mse = numpy.divide(sad, imSz)
            rmse = numpy.sqrt(mse)
            rmseVals.append(rmse)

        return rmseVals
//  [python]
def test_write_float_surface(self):
        width, height, depth = self.dimensions
        dim = 3 if depth != 0 else 2 if height != 0 else 1

        # generate input data and allocate output buffer
        shape = (depth, height, width) if dim == 3 else \
                (height, width) if dim == 2 else \
                (width,)

        # prepare input, output, and surface memory
        real_output = cupy.zeros(shape, dtype=cupy.float32)
        assert real_output.flags['C_CONTIGUOUS']
        ch = ChannelFormatDescriptor(32, 0, 0, 0,
                                     runtime.cudaChannelFormatKindFloat)
        expected_output = cupy.arange(numpy.prod(shape), dtype=cupy.float32)
        expected_output = expected_output.reshape(shape) * 3.0
        assert expected_output.flags['C_CONTIGUOUS']

        # create resource descriptor
        # note that surface memory only support CUDA array
        arr = CUDAarray(ch, width, height, depth,
                        runtime.cudaArraySurfaceLoadStore)
        arr.copy_from(real_output)  # init to zero
        res = ResourceDescriptor(runtime.cudaResourceTypeArray, cuArr=arr)

        # create a surface object; currently we don't support surface reference
        surfobj = SurfaceObject(res)
        mod = cupy.RawModule(code=source_surfobj)

        # get and launch the kernel
        ker_name = 'writeKernel'
        ker_name += '3D' if dim == 3 else '2D' if dim == 2 else '1D'
        ker = mod.get_function(ker_name)
        block = (4, 4, 2) if dim == 3 else (4, 4) if dim == 2 else (4,)
        grid = ()
        args = (surfobj,)
        if dim >= 1:
            grid_x = (width + block[0] - 1)//block[0]
            grid = grid + (grid_x,)
            args = args + (width,)
        if dim >= 2:
            grid_y = (height + block[1] - 1)//block[1]
            grid = grid + (grid_y,)
            args = args + (height,)
        if dim == 3:
            grid_z = (depth + block[2] - 1)//block[2]
            grid = grid + (grid_z,)
            args = args + (depth,)
        ker(grid, block, args)

        # validate result
        arr.copy_to(real_output)
        assert (real_output == expected_output).all()
//  [python]
def test_convolutional():
    x = tensor.tensor4('x')
    num_channels = 4
    num_filters = 3
    batch_size = 5
    filter_size = (3, 3)
    conv = Convolutional(filter_size, num_filters, num_channels,
                         image_size=(17, 13), weights_init=Constant(1.),
                         biases_init=Constant(5.))
    conv.initialize()
    y = conv.apply(x)
    func = function([x], y)

    x_val = numpy.ones((batch_size, num_channels, 17, 13),
                       dtype=theano.config.floatX)
    assert_allclose(func(x_val),
                    n.prod(filter_size) * num_channels *
                    numpy.ones((batch_size, num_filters, 15, 11)) + 5)
    conv.image_size = (17, 13)
    conv.batch_size = 2  # This should have effect on get_dim
    assert conv.get_dim('output') == (num_filters, 15, 11)
//  [python]
def test_convolutional():
    x = tensor.tensor4('x')
    num_channels = 4
    num_filters = 3
    batch_size = 5
    filter_size = (3, 3)
    conv = Convolutional(filter_size, num_filters, num_channels,
                         image_size=(17, 13), weights_init=Constant(1.),
                         biases_init=Constant(5.))
    conv.initialize()
    y = conv.apply(x)
    func = function([x], y)

    x_val = numpy.ones((batch_size, num_channels, 17, 13),
                       dtype=theano.config.floatX)
    assert_allclose(func(x_val),
                    numpy.prod(filter_size) * num_channels *
                    numpy.ones((batch_size, num_filters, 15, 11)) + 5)
    conv.image_size = (17, 13)
    conv.batch_size = 2  # This should have effect on get_dim
    assert conv.get_dim('output') == (num_filters, 15, 11)
//  [python]
def _transform_general(self, func, *args, **kwargs):
        from pandas.tools.merge import concat

        applied = []
        obj = self._obj_with_exclusions
        gen = self.grouper.get_iterator(obj, axis=self.axis)
        fast_path, slow_path = self._define_paths(func, *args, **kwargs)

        path = None
        for name, group in gen:
            object.__setattr__(group, 'name', name)

            if path is None:
                # Try slow path and fast path.
                try:
                    path, res = self._choose_path(fast_path, slow_path, group)
                except TypeError:
                    return self._transform_item_by_item(obj, fast_path)
                except ValueError:
                    msg = 'transform must return a scalar value for each group'
                    raise ValueError(msg)
            else:
                res = path(group)

            if isinstance(res, Series):

                # we need to broadcast across the
                # other dimension; this will preserve dtypes
                # GH14457
                if not np.prod(group.shape):
                    continue
                elif res.index.is_(obj.index):
                    r = concat([res] * len(group.columns), axis=1)
                    r.columns = group.columns
                    r.index = group.index
                else:
                    r = DataFrame(
                        np.concatenate([res.values] * len(group.index)
                                       ).reshape(group.shape),
                        columns=group.columns, index=group.index)

                applied.append(r)
            else:
                applied.append(res)

        concat_index = obj.columns if self.axis == 0 else obj.index
        concatenated = concat(applied, join_axes=[concat_index],
                              axis=self.axis, verify_integrity=False)
        return self._set_result_index_ordered(concatenated)
//  [python]
def test_tied_biases():
    x = tensor.tensor4('x')
    num_channels = 4
    num_filters = 3
    batch_size = 5
    filter_size = (3, 3)
    conv = Convolutional(filter_size, num_filters, num_channels,
                         weights_init=Constant(1.), biases_init=Constant(2.),
                         tied_biases=True)
    conv.initialize()
    y = conv.apply(x)
    func = function([x], y)

    # Tied biases allows to pass images of different sizes
    x_val_1 = numpy.ones((batch_size, num_channels, 10,
                          12), dtype=theano.config.floatX)
    x_val_2 = numpy.ones((batch_size, num_channels, 23,
                          19), dtype=theano.config.floatX)

    assert_allclose(func(x_val_1),
                    n.prod(filter_size) * num_channels *
                    numpy.ones((batch_size, num_filters, 8, 10)) + 2)
    assert_allclose(func(x_val_2),
                    n.prod(filter_size) * num_channels *
                    numpy.ones((batch_size, num_filters, 21, 17)) + 2)
//  [python]
def test_tied_biases():
    x = tensor.tensor4('x')
    num_channels = 4
    num_filters = 3
    batch_size = 5
    filter_size = (3, 3)
    conv = Convolutional(filter_size, num_filters, num_channels,
                         weights_init=Constant(1.), biases_init=Constant(2.),
                         tied_biases=True)
    conv.initialize()
    y = conv.apply(x)
    func = function([x], y)

    # Tied biases allows to pass images of different sizes
    x_val_1 = numpy.ones((batch_size, num_channels, 10,
                          12), dtype=theano.config.floatX)
    x_val_2 = numpy.ones((batch_size, num_channels, 23,
                          19), dtype=theano.config.floatX)

    assert_allclose(func(x_val_1),
                    numpy.prod(filter_size) * num_channels *
                    numpy.ones((batch_size, num_filters, 8, 10)) + 2)
    assert_allclose(func(x_val_2),
                    numpy.prod(filter_size) * num_channels *
                    numpy.ones((batch_size, num_filters, 21, 17)) + 2)
//  [python]
def op_numpy(self, matrix, dims, sel):
        sel.sort()
        ndims = len(dims)
        dkeep = [dims[x] for x in sel]
        qtrace = list(set(range(ndims)) - set(sel))
        dtrace = [dims[x] for x in qtrace]

        matrix = matrix.reshape(dims + dims)
        matrix = matrix.transpose(
            qtrace + [ndims + i for i in qtrace] + sel + [ndims + i for i in sel]
        )
        matrix = matrix.reshape(
            [
                n.prod(dtrace, dtype=int),
                n.prod(dtrace, dtype=int),
                n.prod(dkeep, dtype=int),
                n.prod(dkeep, dtype=int),
            ]
        )
        return np.trace(matrix)
//  [python]
def op_numpy(self, matrix, dims, sel):
        sel.sort()
        ndims = len(dims)
        dkeep = [dims[x] for x in sel]
        qtrace = list(set(range(ndims)) - set(sel))
        dtrace = [dims[x] for x in qtrace]

        matrix = matrix.reshape(dims + dims)
        matrix = matrix.transpose(
            qtrace + [ndims + i for i in qtrace] + sel + [ndims + i for i in sel]
        )
        matrix = matrix.reshape(
            [
                np.prod(dtrace, dtype=int),
                np.prod(dtrace, dtype=int),
                np.prod(dkeep, dtype=int),
                np.prod(dkeep, dtype=int),
            ]
        )
        return np.trace(matrix)
//  [python]
def setUp(self):
        # x, W, and b are set so that the result of forward
        # propagation gets stable, meaning that their small perturbations
        # do not change :math:`argmax_{j} W_{ij\cdot} x + b_{ij}`.

        x_shape = (self.batchsize, ) + self.in_shape
        self.x = numpy.random.uniform(
            -0.05, 0.05, x_shape).astype(numpy.float32) + 1
        self.gy = numpy.random.uniform(
            -0.05, 0.05, (self.batchsize, self.out_size)
        ).astype(numpy.float32)

        in_size = numpy.prod(self.in_shape)
        initialW = numpy.random.uniform(
            -0.05, 0.05, (self.out_size, self.pool_size, in_size)
        ).astype(numpy.float32)
        for o in six.moves.range(self.out_size):
            w = numpy.arange(in_size, dtype=numpy.float32) + 1
            for c in six.moves.range(self.pool_size):
                initialW[o, c, :] += w * c

        if self.initial_bias == 'random':
            initial_bias = numpy.random.uniform(
                -0.05, 0.05, (self.out_size, self.pool_size))
        elif self.initial_bias == 'scalar':
            initial_bias = numpy.full(
                (self.out_size, self.pool_size), 5, dtype=numpy.float32)
        elif self.initial_bias is None:
            initial_bias = None

        self.link = links.Maxout(in_size, self.out_size, self.pool_size,
                                 initialW, initial_bias)

        self.y = _maxout(self.x, initialW, initial_bias)
        self.link.cleargrads()
//  [python]
def test_sparse_nd():
    rng = numpy.random.RandomState(1)

    def check_sparse(rng, axis, num_init, shape, weights_init=Constant(1.)):
        weights = SparseND(axis=axis, num_init=num_init,
                           weights_init=weights_init).generate(rng, shape)
        assert weights.shape == shape
        assert weights.dtype == theano.config.floatX
        if isinstance(num_init, numbers.Integral):
            nnz = n.prod([s for i, s in enumerate(shape)
                              if i in pack(axis)]) * num_init
            assert numpy.count_nonzero(weights) == nnz
        else:
            atom_size = n.prod([s for i, s in enumerate(shape)
                                    if i not in pack(axis)])
            nnz_atom = int(num_init * atom_size)
            num_atoms = n.prod([s for i, s in enumerate(shape)
                                    if i in pack(axis)])
            nnz = nnz_atom * num_atoms
            assert numpy.count_nonzero(weights) == nnz

    yield check_sparse, rng, 1, 5, (10, 11)
    yield check_sparse, rng, 2, 3, (7, 8, 9)
    yield check_sparse, rng, (2, 3), 5. / 6., (2, 3, 5, 7)
    yield check_sparse, rng, (0, 1), 3, (3, 5, 7, 11)
    yield check_sparse, rng, (0, 2, 3), 0.5, (2, 3, 2, 6)
//  [python]
def test_sparse_nd():
    rng = numpy.random.RandomState(1)

    def check_sparse(rng, axis, num_init, shape, weights_init=Constant(1.)):
        weights = SparseND(axis=axis, num_init=num_init,
                           weights_init=weights_init).generate(rng, shape)
        assert weights.shape == shape
        assert weights.dtype == theano.config.floatX
        if isinstance(num_init, numbers.Integral):
            nnz = numpy.prod([s for i, s in enumerate(shape)
                              if i in pack(axis)]) * num_init
            assert numpy.count_nonzero(weights) == nnz
        else:
            atom_size = numpy.prod([s for i, s in enumerate(shape)
                                    if i not in pack(axis)])
            nnz_atom = int(num_init * atom_size)
            num_atoms = numpy.prod([s for i, s in enumerate(shape)
                                    if i in pack(axis)])
            nnz = nnz_atom * num_atoms
            assert numpy.count_nonzero(weights) == nnz

    yield check_sparse, rng, 1, 5, (10, 11)
    yield check_sparse, rng, 2, 3, (7, 8, 9)
    yield check_sparse, rng, (2, 3), 5. / 6., (2, 3, 5, 7)
    yield check_sparse, rng, (0, 1), 3, (3, 5, 7, 11)
    yield check_sparse, rng, (0, 2, 3), 0.5, (2, 3, 2, 6)
//  [python]
def normalize_numba(array, min_value, max_value, epsilon):
        for _ in prange(numpy.multiply.reduce(numpy.array(array.shape))):
            array.flat[_] -= min_value
            array.flat[_] /= max_value - min_value + epsilon
//  [python]
def normalize_numba(array, min_value, max_value, epsilon):
        for _ in prange(numpy.prod(numpy.array(array.shape))):
            array.flat[_] -= min_value
            array.flat[_] /= max_value - min_value + epsilon
//  [python]
def gradients(scalar, variables):
    """Compute the gradients of a scalar w.r.t to a given list of variables.

    Arguments
    ---------
    scalar: :class:`symjax.tensor.base.Tensor`
        the variable to differentiate

    variables: List or Tuple
        the variables used to compute the derivative.

    Returns
    -------

        gradients: Tuple
            the sequency of gradients ordered as given in the input variables

    Example
    -------

    .. doctest::

        >>> import symjax
        >>> w = symjax.tensor.ones(3)
        >>> x = symjax.tensor.Variable(2., name='x', dtype='float32')
        >>> l = (w ** 2).sum() * x
        >>> g = symjax.gradients(l, [w])[0]
        >>> f = symjax.function(outputs=g, updates={x:x + 1})
        >>> for i in range(2):
        ...    print(f())
        [4. 4. 4.]
        [6. 6. 6.]

    """
    if n.prod(scalar.shape) != 1:
        raise RuntimeError("the variable to differentiate is not a scalar")
    if not isinstance(scalar, t.Tensor):
        raise RuntimeError("the variable used in gradients should be a Tensor type")

    if isinstance(variables, t.Tensor):
        input_variables = [variables]
        input_list = False
    else:
        input_variables = variables.copy()
        input_list = True

    # get the argnum of the variables that we differentiate one
    argnums = list(range(len(input_variables)))

    # get all the roots of the scalar, this is needed as otherwise they are not
    # as the input of the gradient function and thus a change of
    # their value will not change the gradient computation, we also ensure
    # uniqueness
    input_variables += [
        i for i in current_graph().roots(scalar) if i not in input_variables
    ]

    # create a dummy function that is needed for jax to compute a gradient func
    # this function is the one that builds the graph of computation from all
    # roots
    # to the scalar varible s.t. automatic diffenrentiation can be applied

    def internal_gradient_function(*args):
        return current_graph().get(
            scalar,
            {input_variables[i]: args[i] for i in range(len(input_variables))},
        )

    # now we obtain the grad function. In fact, Jax returns a function that,
    # when it is called, returns the gradient values, this function is then
    # used to generate the Tuple of symbolic variables
    grad_fn = jax.grad(internal_gradient_function, argnums)
    wrap_fn = t.jax_wrap(grad_fn)
    if input_list:
        return wrap_fn(*input_variables)
    else:
        return wrap_fn(*input_variables)[0]
//  [python]
def gradients(scalar, variables):
    """Compute the gradients of a scalar w.r.t to a given list of variables.

    Arguments
    ---------
    scalar: :class:`symjax.tensor.base.Tensor`
        the variable to differentiate

    variables: List or Tuple
        the variables used to compute the derivative.

    Returns
    -------

        gradients: Tuple
            the sequency of gradients ordered as given in the input variables

    Example
    -------

    .. doctest::

        >>> import symjax
        >>> w = symjax.tensor.ones(3)
        >>> x = symjax.tensor.Variable(2., name='x', dtype='float32')
        >>> l = (w ** 2).sum() * x
        >>> g = symjax.gradients(l, [w])[0]
        >>> f = symjax.function(outputs=g, updates={x:x + 1})
        >>> for i in range(2):
        ...    print(f())
        [4. 4. 4.]
        [6. 6. 6.]

    """
    if numpy.prod(scalar.shape) != 1:
        raise RuntimeError("the variable to differentiate is not a scalar")
    if not isinstance(scalar, t.Tensor):
        raise RuntimeError("the variable used in gradients should be a Tensor type")

    if isinstance(variables, t.Tensor):
        input_variables = [variables]
        input_list = False
    else:
        input_variables = variables.copy()
        input_list = True

    # get the argnum of the variables that we differentiate one
    argnums = list(range(len(input_variables)))

    # get all the roots of the scalar, this is needed as otherwise they are not
    # as the input of the gradient function and thus a change of
    # their value will not change the gradient computation, we also ensure
    # uniqueness
    input_variables += [
        i for i in current_graph().roots(scalar) if i not in input_variables
    ]

    # create a dummy function that is needed for jax to compute a gradient func
    # this function is the one that builds the graph of computation from all
    # roots
    # to the scalar varible s.t. automatic diffenrentiation can be applied

    def internal_gradient_function(*args):
        return current_graph().get(
            scalar,
            {input_variables[i]: args[i] for i in range(len(input_variables))},
        )

    # now we obtain the grad function. In fact, Jax returns a function that,
    # when it is called, returns the gradient values, this function is then
    # used to generate the Tuple of symbolic variables
    grad_fn = jax.grad(internal_gradient_function, argnums)
    wrap_fn = t.jax_wrap(grad_fn)
    if input_list:
        return wrap_fn(*input_variables)
    else:
        return wrap_fn(*input_variables)[0]
//  [python]
def interpolate(batch_size, noise):
    half_batch = batch_size//2
    z1 = noise[:half_batch]
    z2 = noise[half_batch:]

    generate = []
    for idx, ratio in enumerate(np.linspace(0, 1, 10)):
        z = np.stack([slerp(ratio, r1, r2) for r1, r2 in zip(z1, z2)])
        generate.append(G(z))
    generate = np.stack(generate).transpose([1,0,2,3,4])
    for idx, img in enumerate(generate):
        save_image(img, os.path.join(os.getcwd()+'/interpolate', 'interp_G_{}.png'.format(idx)), nrow=10)

    all_img_num = np.multiply.reduce(generate.shape[:2])
    new_shape = [all_img_num] + list(generated.shape[2:])
    batch_generate = np.reshape(generate, new_shape)
    save_image(batch_generate, os.path.join(os.getcwd()+'/interpolate', 'interp_G.png'), nrow=10)
//  [python]
def interpolate(batch_size, noise):
    half_batch = batch_size//2
    z1 = noise[:half_batch]
    z2 = noise[half_batch:]

    generate = []
    for idx, ratio in enumerate(np.linspace(0, 1, 10)):
        z = np.stack([slerp(ratio, r1, r2) for r1, r2 in zip(z1, z2)])
        generate.append(G(z))
    generate = np.stack(generate).transpose([1,0,2,3,4])
    for idx, img in enumerate(generate):
        save_image(img, os.path.join(os.getcwd()+'/interpolate', 'interp_G_{}.png'.format(idx)), nrow=10)

    all_img_num = np.prod(generate.shape[:2])
    new_shape = [all_img_num] + list(generated.shape[2:])
    batch_generate = np.reshape(generate, new_shape)
    save_image(batch_generate, os.path.join(os.getcwd()+'/interpolate', 'interp_G.png'), nrow=10)
//  [python]
def denormalize_numba(array, min_value, max_value, epsilon):
        for _ in prange(numpy.multiply.reduce(numpy.array(array.shape))):
            array.flat[_] *= max_value - min_value + epsilon
            array.flat[_] += min_value
//  [python]
def denormalize_numba(array, min_value, max_value, epsilon):
        for _ in prange(numpy.prod(numpy.array(array.shape))):
            array.flat[_] *= max_value - min_value + epsilon
            array.flat[_] += min_value
//  [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 choice(self, a, size=None, replace=True, p=None):
        """Returns an array of random values from a given 1-D array.

        .. seealso::
            - :func:`cupy.random.choice` for full documentation
            - :meth:`numpy.random.choice`

        """
        if a is None:
            raise ValueError('a must be 1-dimensional or an integer')
        if isinstance(a, cupy.ndarray) and a.ndim == 0:
            raise NotImplementedError
        if isinstance(a, int):
            a_size = a
            if a_size < 0:
                raise ValueError('a must be greater than or equal to 0')
        else:
            a = cupy.array(a, copy=False)
            if a.ndim != 1:
                raise ValueError('a must be 1-dimensional or an integer')
            a_size = len(a)

        if p is not None:
            p = cupy.array(p)
            if p.ndim != 1:
                raise ValueError('p must be 1-dimensional')
            if len(p) != a_size:
                raise ValueError('a and p must have same size')
            if not (p >= 0).all():
                raise ValueError('probabilities are not non-negative')
            p_sum = cupy.sum(p).get()
            if not numpy.allclose(p_sum, 1):
                raise ValueError('probabilities do not sum to 1')

        if size is None:
            raise NotImplementedError(
                'choice() without specifying size is not supported yet')
        shape = size
        size = numpy.prod(shape)

        if a_size == 0 and size > 0:
            raise ValueError('a cannot be empty unless no samples are taken')

        if not replace and p is None:
            if a_size < size:
                raise ValueError(
                    'Cannot take a larger sample than population when '
                    '\'replace=False\'')
            if isinstance(a, int):
                indices = cupy.arange(a, dtype='l')
            else:
                indices = a.copy()
            self.shuffle(indices)
            return indices[:size].reshape(shape)

        if not replace:
            raise NotImplementedError

        if p is not None:
            p = cupy.broadcast_to(p, (size, a_size))
            index = cupy.argmax(cupy.log(p) +
                                self.gumbel(size=(size, a_size)),
                                axis=1)
            if not isinstance(shape, int):
                index = cupy.reshape(index, shape)
        else:
            if a_size == 0:  # TODO: (#4511) Fix `randint` instead
                a_size = 1
            index = self.randint(0, a_size, size=shape)
            # Align the dtype with NumPy
            index = index.astype(cupy.int64, copy=False)

        if isinstance(a, int):
            return index

        if index.ndim == 0:
            return cupy.array(a[index], dtype=a.dtype)

        return a[index]
//  [python]
def _get_qkv_length(
        self,
        q_shape: List[int],
        k_shape: List[int],
        v_shape: List[int],
    ) -> Tuple[int]:
        q_N = n.prod(q_shape) + 1 if self.has_cls_embed else n.prod(q_shape)
        k_N = n.prod(k_shape) + 1 if self.has_cls_embed else n.prod(k_shape)
        v_N = n.prod(v_shape) + 1 if self.has_cls_embed else n.prod(v_shape)
        return q_N, k_N, v_N
//  [python]
def _get_qkv_length(
        self,
        q_shape: List[int],
        k_shape: List[int],
        v_shape: List[int],
    ) -> Tuple[int]:
        q_N = numpy.prod(q_shape) + 1 if self.has_cls_embed else numpy.prod(q_shape)
        k_N = numpy.prod(k_shape) + 1 if self.has_cls_embed else numpy.prod(k_shape)
        v_N = numpy.prod(v_shape) + 1 if self.has_cls_embed else numpy.prod(v_shape)
        return q_N, k_N, v_N
//  [python]
def get_signature_total_size(input_signature: SignatureType) -> int:
  """Gets the total number of elements in a single problem instance

  Args:
    input_signature: An input signature to calculate the number of elements in
  """
  total_size = 0
  for input_key in input_signature:
    total_size += n.prod(input_signature[input_key][0])
  return total_size
//  [python]
def get_signature_total_size(input_signature: SignatureType) -> int:
  """Gets the total number of elements in a single problem instance

  Args:
    input_signature: An input signature to calculate the number of elements in
  """
  total_size = 0
  for input_key in input_signature:
    total_size += numpy.prod(input_signature[input_key][0])
  return total_size
//  [python]
def _pseudo_inverse_dense(L, rhoss, method='direct'):
    """
    Internal function for computing the pseudo inverse of an Liouvillian using
    dense matrix methods. See pseudo_inverse for details.
    """
    if method == 'direct':
        rho_vec = np.transpose(mat2vec(rhoss.full()))

        tr_mat = tensor([identity(n) for n in L.dims[0][0]])
        tr_vec = np.transpose(mat2vec(tr_mat.full()))

        N = n.prod(L.dims[0][0])
        I = np.identity(N * N)
        P = np.kron(np.transpose(rho_vec), tr_vec)
        Q = I - P
        LIQ = np.linalg.solve(L.full(), Q)
        R = np.dot(Q, LIQ)

        return Qobj(R, dims=L.dims)

    elif method == 'numpy':
        return Qobj(np.linalg.pinv(L.full()), dims=L.dims)

    elif method == 'scipy':
        return Qobj(la.pinv(L.full()), dims=L.dims)

    elif method == 'scipy2':
        return Qobj(la.pinv2(L.full()), dims=L.dims)

    else:
        raise ValueError("Unsupported method '%s'. Use 'direct' or 'numpy'" %
                         method)
//  [python]
def _pseudo_inverse_dense(L, rhoss, method='direct'):
    """
    Internal function for computing the pseudo inverse of an Liouvillian using
    dense matrix methods. See pseudo_inverse for details.
    """
    if method == 'direct':
        rho_vec = np.transpose(mat2vec(rhoss.full()))

        tr_mat = tensor([identity(n) for n in L.dims[0][0]])
        tr_vec = np.transpose(mat2vec(tr_mat.full()))

        N = np.prod(L.dims[0][0])
        I = np.identity(N * N)
        P = np.kron(np.transpose(rho_vec), tr_vec)
        Q = I - P
        LIQ = np.linalg.solve(L.full(), Q)
        R = np.dot(Q, LIQ)

        return Qobj(R, dims=L.dims)

    elif method == 'numpy':
        return Qobj(np.linalg.pinv(L.full()), dims=L.dims)

    elif method == 'scipy':
        return Qobj(la.pinv(L.full()), dims=L.dims)

    elif method == 'scipy2':
        return Qobj(la.pinv2(L.full()), dims=L.dims)

    else:
        raise ValueError("Unsupported method '%s'. Use 'direct' or 'numpy'" %
                         method)
//  [python]
def get_max_part_size(input_signature: SignatureType) -> int:
  """Gets the size (as a single scalar) of the largest feature in terms of
  the number of elements.

  Args:
    input_signature: The input signature that we want to find the largest
      feature in.
  """
  part_sizes = numpy.empty(len(input_signature))
  for index, input_key in enumerate(input_signature):
    part_sizes[index] = numpy.prod(input_signature[input_key][0])
  return numpy.max(part_sizes)
//  [python]
def pack_flat_array_into_input(
    flat_array: numpy.typing.ArrayLike,
    signature_spec: SignatureType) -> types.NestedTensorSpec:
  """Packs a flat array into a nested tensor spec to feed into a model

  Args:
    flat_array: The data to be packed back into the specified nested tensor
      specification
    signature_spec: A signature that is used to create the correct structure
      for all of the values in the flat array
  """
  output_input_dict = {}
  current_index = 0
  for needed_input in signature_spec:
    part_size = n.prod(signature_spec[needed_input][0])
    needed_subset = flat_array[current_index:current_index + part_size]
    current_index += part_size
    output_input_dict[needed_input] = tf.cast(
        tf.constant(needed_subset, shape=signature_spec[needed_input][0]),
        dtype=signature_spec[needed_input][1])
  return output_input_dict
//  [python]
def pack_flat_array_into_input(
    flat_array: numpy.typing.ArrayLike,
    signature_spec: SignatureType) -> types.NestedTensorSpec:
  """Packs a flat array into a nested tensor spec to feed into a model

  Args:
    flat_array: The data to be packed back into the specified nested tensor
      specification
    signature_spec: A signature that is used to create the correct structure
      for all of the values in the flat array
  """
  output_input_dict = {}
  current_index = 0
  for needed_input in signature_spec:
    part_size = numpy.prod(signature_spec[needed_input][0])
    needed_subset = flat_array[current_index:current_index + part_size]
    current_index += part_size
    output_input_dict[needed_input] = tf.cast(
        tf.constant(needed_subset, shape=signature_spec[needed_input][0]),
        dtype=signature_spec[needed_input][1])
  return output_input_dict
//  [python]
def collapse_values(input_signature: SignatureType,
                    shap_values: numpy.typing.ArrayLike,
                    num_examples: int) -> numpy.typing.ArrayLike:
  """Collapses shap values so that there is only a single value per feature

  Args:
    input_signature: The signature of the model input. Used to determine what
      (if any) features need to be collapsed.
    shap_values: A numpy array of shap values that need to be processed.
  """
  output_shap_values = numpy.empty((num_examples, len(input_signature)))
  for i in range(0, num_examples):
    current_index = 0
    current_feature = 0
    for input_key in input_signature:
      part_size = n.prod(input_signature[input_key][0])
      output_shap_values[i, current_feature] = numpy.sum(
          shap_values[i, current_index:current_index + part_size])
      current_feature += 1
      current_index += part_size
  return output_shap_values
//  [python]
def collapse_values(input_signature: SignatureType,
                    shap_values: numpy.typing.ArrayLike,
                    num_examples: int) -> numpy.typing.ArrayLike:
  """Collapses shap values so that there is only a single value per feature

  Args:
    input_signature: The signature of the model input. Used to determine what
      (if any) features need to be collapsed.
    shap_values: A numpy array of shap values that need to be processed.
  """
  output_shap_values = numpy.empty((num_examples, len(input_signature)))
  for i in range(0, num_examples):
    current_index = 0
    current_feature = 0
    for input_key in input_signature:
      part_size = numpy.prod(input_signature[input_key][0])
      output_shap_values[i, current_feature] = numpy.sum(
          shap_values[i, current_index:current_index + part_size])
      current_feature += 1
      current_index += part_size
  return output_shap_values
//  [python]
def _pseudo_inverse_sparse(L, rhoss, method='splu', use_umfpack=False,
                           use_rcm=False):
    """
    Internal function for computing the pseudo inverse of an Liouvillian using
    sparse matrix methods. See pseudo_inverse for details.
    """

    N = n.prod(L.dims[0][0])

    rhoss_vec = operator_to_vector(rhoss)

    tr_op = tensor([identity(n) for n in L.dims[0][0]])
    tr_op_vec = operator_to_vector(tr_op)

    P = sp.kron(rhoss_vec.data, tr_op_vec.data.T, format='csc')
    I = sp.eye(N*N, N*N, format='csc')
    Q = I - P

    if use_rcm:
        perm = reverse_cuthill_mckee(L.data)
        A = sp_permute(L.data, perm, perm, 'csc').tocsc()
        Q = sp_permute(Q, perm, perm, 'csc')
        permc_spec = 'NATURAL'
    else:
        A = L.data.tocsc()
        A.sort_indices()
        permc_spec = 'COLAMD'

    if method == 'spsolve':
        sp.linalg.use_solver(assumeSortedIndices=True, useUmfpack=use_umfpack)
        LIQ = sp.linalg.spsolve(A, Q)

    elif method == 'splu':
        lu = sp.linalg.splu(A, permc_spec=permc_spec)
        LIQ = lu.solve(Q.toarray())

    elif method == 'spilu':
        lu = sp.linalg.spilu(A, permc_spec=permc_spec,
                             fill_factor=10, drop_tol=1e-8)
        LIQ = lu.solve(Q.toarray())

    else:
        raise ValueError("unsupported method '%s'" % method)

    R = sp.csc_matrix(Q * LIQ)

    if use_rcm:
        rev_perm = np.argsort(perm)
        R = sp_permute(R, rev_perm, rev_perm, 'csc')

    return Qobj(R, dims=L.dims)
//  [python]
def _pseudo_inverse_sparse(L, rhoss, method='splu', use_umfpack=False,
                           use_rcm=False):
    """
    Internal function for computing the pseudo inverse of an Liouvillian using
    sparse matrix methods. See pseudo_inverse for details.
    """

    N = np.prod(L.dims[0][0])

    rhoss_vec = operator_to_vector(rhoss)

    tr_op = tensor([identity(n) for n in L.dims[0][0]])
    tr_op_vec = operator_to_vector(tr_op)

    P = sp.kron(rhoss_vec.data, tr_op_vec.data.T, format='csc')
    I = sp.eye(N*N, N*N, format='csc')
    Q = I - P

    if use_rcm:
        perm = reverse_cuthill_mckee(L.data)
        A = sp_permute(L.data, perm, perm, 'csc').tocsc()
        Q = sp_permute(Q, perm, perm, 'csc')
        permc_spec = 'NATURAL'
    else:
        A = L.data.tocsc()
        A.sort_indices()
        permc_spec = 'COLAMD'

    if method == 'spsolve':
        sp.linalg.use_solver(assumeSortedIndices=True, useUmfpack=use_umfpack)
        LIQ = sp.linalg.spsolve(A, Q)

    elif method == 'splu':
        lu = sp.linalg.splu(A, permc_spec=permc_spec)
        LIQ = lu.solve(Q.toarray())

    elif method == 'spilu':
        lu = sp.linalg.spilu(A, permc_spec=permc_spec,
                             fill_factor=10, drop_tol=1e-8)
        LIQ = lu.solve(Q.toarray())

    else:
        raise ValueError("unsupported method '%s'" % method)

    R = sp.csc_matrix(Q * LIQ)

    if use_rcm:
        rev_perm = np.argsort(perm)
        R = sp_permute(R, rev_perm, rev_perm, 'csc')

    return Qobj(R, dims=L.dims)
//  [python]
def qr(
    np,  # TODO: Typing
    tensor: Tensor,
    pivot_axis: int,
    non_negative_diagonal: bool
) -> Tuple[Tensor, Tensor]:
  """Computes the QR decomposition of a tensor.

  See tensornetwork.backends.tensorflow.decompositions for details.
  """
  left_dims = tensor.shape[:pivot_axis]
  right_dims = tensor.shape[pivot_axis:]
  tensor = np.reshape(tensor, [np.prod(left_dims), np.prod(right_dims)])
  q, r = np.linalg.qr(tensor)
  if non_negative_diagonal:
    phases = np.sign(np.diagonal(r))
    q = q * phases
    r = phases.conj()[:, None] * r
  center_dim = q.shape[1]
  q = np.reshape(q, list(left_dims) + [center_dim])
  r = np.reshape(r, [center_dim] + list(right_dims))
  return q, r
//  [python]
def qr(
    np,  # TODO: Typing
    tensor: Tensor,
    pivot_axis: int,
    non_negative_diagonal: bool
) -> Tuple[Tensor, Tensor]:
  """Computes the QR decomposition of a tensor.

  See tensornetwork.backends.tensorflow.decompositions for details.
  """
  left_dims = tensor.shape[:pivot_axis]
  right_dims = tensor.shape[pivot_axis:]
  tensor = np.reshape(tensor, [numpy.prod(left_dims), numpy.prod(right_dims)])
  q, r = np.linalg.qr(tensor)
  if non_negative_diagonal:
    phases = np.sign(np.diagonal(r))
    q = q * phases
    r = phases.conj()[:, None] * r
  center_dim = q.shape[1]
  q = np.reshape(q, list(left_dims) + [center_dim])
  r = np.reshape(r, [center_dim] + list(right_dims))
  return q, r
//  [python]
def product(self, tensor_in: Tensor[T], axis: Shape | None = None) -> ArrayLike:
        return n.prod(tensor_in, axis=axis)
//  [python]
def product(self, tensor_in: Tensor[T], axis: Shape | None = None) -> ArrayLike:
        return np.prod(tensor_in, axis=axis)
//  [python]
def rq(
    np,  # TODO: Typing
    tensor: Tensor,
    pivot_axis: int,
    non_negative_diagonal: bool
) -> Tuple[Tensor, Tensor]:
  """Computes the RQ (reversed QR) decomposition of a tensor.

  See tensornetwork.backends.tensorflow.decompositions for details.
  """
  left_dims = tensor.shape[:pivot_axis]
  right_dims = tensor.shape[pivot_axis:]
  tensor = np.reshape(tensor, [np.prod(left_dims), np.prod(right_dims)])
  q, r = np.linalg.qr(np.conj(np.transpose(tensor)))
  if non_negative_diagonal:
    phases = np.sign(np.diagonal(r))
    q = q * phases
    r = phases.conj()[:, None] * r
  r, q = np.conj(np.transpose(r)), np.conj(
      np.transpose(q))  #M=r*q at this point
  center_dim = r.shape[1]
  r = np.reshape(r, list(left_dims) + [center_dim])
  q = np.reshape(q, [center_dim] + list(right_dims))
  return r, q
//  [python]
def rq(
    np,  # TODO: Typing
    tensor: Tensor,
    pivot_axis: int,
    non_negative_diagonal: bool
) -> Tuple[Tensor, Tensor]:
  """Computes the RQ (reversed QR) decomposition of a tensor.

  See tensornetwork.backends.tensorflow.decompositions for details.
  """
  left_dims = tensor.shape[:pivot_axis]
  right_dims = tensor.shape[pivot_axis:]
  tensor = np.reshape(tensor, [numpy.prod(left_dims), numpy.prod(right_dims)])
  q, r = np.linalg.qr(np.conj(np.transpose(tensor)))
  if non_negative_diagonal:
    phases = np.sign(np.diagonal(r))
    q = q * phases
    r = phases.conj()[:, None] * r
  r, q = np.conj(np.transpose(r)), np.conj(
      np.transpose(q))  #M=r*q at this point
  center_dim = r.shape[1]
  r = np.reshape(r, list(left_dims) + [center_dim])
  q = np.reshape(q, [center_dim] + list(right_dims))
  return r, q
//  [python]
dst_mask_values=None
if domask:
    dst_mask_values=numpy.array([0])

regrid = esmpy.Regrid(srcfield, dstfield,
                     regrid_method=esmpy.RegridMethod.BILINEAR,
                     unmapped_action=esmpy.UnmappedAction.ERROR,
                     dst_mask_values=dst_mask_values)

# do the regridding from source to destination field
dstfield = regrid(srcfield, dstfield, zero_region=esmpy.Region.SELECT)

# compute the mean relative error
num_nodes = n.prod(xctfield.data.shape[:])
//  [python]
dst_mask_values=None
if domask:
    dst_mask_values=numpy.array([0])

regrid = esmpy.Regrid(srcfield, dstfield,
                     regrid_method=esmpy.RegridMethod.BILINEAR,
                     unmapped_action=esmpy.UnmappedAction.ERROR,
                     dst_mask_values=dst_mask_values)

# do the regridding from source to destination field
dstfield = regrid(srcfield, dstfield, zero_region=esmpy.Region.SELECT)

# compute the mean relative error
num_nodes = numpy.prod(xctfield.data.shape[:])
//  [python]
def batch_jacobian(outputs, inputs, create_graph=True):
    """Computes the jacobian of outputs with respect to inputs, assuming the first dimension of both are the minibatch.

    Based on gelijergensen's code at https://gist.github.com/sbarratt/37356c46ad1350d4c30aefbd488a4faa.

    :param outputs: tensor for the output of some function
    :param inputs: tensor for the input of some function (probably a vector)
    :param create_graph: set True for the resulting jacobian to be differentible
    :returns: a tensor of size (outputs.size() + inputs.size()) containing the
        jacobian of outputs with respect to inputs
    """

    jac = calculate_jacobian(outputs, inputs)
    jac = jac.view((outputs.size(0), n.prod(outputs.size()[1:]), inputs.size(0), n.prod(inputs.size()[1:])))
    jac = torch.einsum("bibj->bij", jac)

    if create_graph:
        jac.requires_grad_()

    return jac
//  [python]
def batch_jacobian(outputs, inputs, create_graph=True):
    """Computes the jacobian of outputs with respect to inputs, assuming the first dimension of both are the minibatch.

    Based on gelijergensen's code at https://gist.github.com/sbarratt/37356c46ad1350d4c30aefbd488a4faa.

    :param outputs: tensor for the output of some function
    :param inputs: tensor for the input of some function (probably a vector)
    :param create_graph: set True for the resulting jacobian to be differentible
    :returns: a tensor of size (outputs.size() + inputs.size()) containing the
        jacobian of outputs with respect to inputs
    """

    jac = calculate_jacobian(outputs, inputs)
    jac = jac.view((outputs.size(0), np.prod(outputs.size()[1:]), inputs.size(0), np.prod(inputs.size()[1:])))
    jac = torch.einsum("bibj->bij", jac)

    if create_graph:
        jac.requires_grad_()

    return jac