jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def check_shape(args, current_shape=None):
    """Check validity of the shape"""

    if len(args) == 0:
        raise TypeError("function missing 1 required positional argument: "
                        "'shape'")

    elif len(args) == 1:
        try:
            shape_iter = iter(args[0])
        except TypeError:
            new_shape = (operator.index(args[0]), )
        else:
            new_shape = tuple(operator.index(arg) for arg in shape_iter)
    else:
        new_shape = tuple(operator.index(arg) for arg in args)

    if current_shape is None:
        if len(new_shape) != 2:
            raise ValueError('shape must be a 2-tuple of positive integers')
        elif new_shape[0] < 0 or new_shape[1] < 0:
            raise ValueError("'shape' elements cannot be negative")

    else:
        current_size = reduce(lambda x, y: x * y, current_shape)

        negative_indexes = [i for i, x in enumerate(new_shape) if x < 0]
        if len(negative_indexes) == 0:
            new_size = reduce(lambda x, y: x * y, new_shape)
            if new_size != current_size:
                raise ValueError('cannot reshape array of size {} into shape'
                                 '{}'.format(current_size, new_shape))
        elif len(negative_indexes) == 1:
            skip = negative_indexes[0]
            specified = reduce(lambda x, y: x * y, new_shape[0:skip] + new_shape[skip+1:])
            unspecified, remainder = divmod(current_size, specified)
            if remainder != 0:
                err_shape = tuple('newshape'if x < 0 else x for x in new_shape)
                raise ValueError('cannot reshape array of size {} into shape'
                                 '{}'.format(current_size, err_shape))
            new_shape = new_shape[0:skip] + (unspecified,) + new_shape[skip+1:]
        else:
            raise ValueError('can only specify one unknown dimension')

    if len(new_shape) != 2:
        raise ValueError('matrix shape must be two-dimensional')

    return new_shape
//  [python]
def check_shape(args, current_shape=None):
    """Check validity of the shape"""

    if len(args) == 0:
        raise TypeError("function missing 1 required positional argument: "
                        "'shape'")

    elif len(args) == 1:
        try:
            shape_iter = iter(args[0])
        except TypeError:
            new_shape = (operator.index(args[0]), )
        else:
            new_shape = tuple(operator.index(arg) for arg in shape_iter)
    else:
        new_shape = tuple(operator.index(arg) for arg in args)

    if current_shape is None:
        if len(new_shape) != 2:
            raise ValueError('shape must be a 2-tuple of positive integers')
        elif new_shape[0] < 0 or new_shape[1] < 0:
            raise ValueError("'shape' elements cannot be negative")

    else:
        current_size = numpy.prod(current_shape)

        negative_indexes = [i for i, x in enumerate(new_shape) if x < 0]
        if len(negative_indexes) == 0:
            new_size = numpy.prod(new_shape)
            if new_size != current_size:
                raise ValueError('cannot reshape array of size {} into shape'
                                 '{}'.format(current_size, new_shape))
        elif len(negative_indexes) == 1:
            skip = negative_indexes[0]
            specified = numpy.prod(new_shape[0:skip] + new_shape[skip+1:])
            unspecified, remainder = divmod(current_size, specified)
            if remainder != 0:
                err_shape = tuple('newshape'if x < 0 else x for x in new_shape)
                raise ValueError('cannot reshape array of size {} into shape'
                                 '{}'.format(current_size, err_shape))
            new_shape = new_shape[0:skip] + (unspecified,) + new_shape[skip+1:]
        else:
            raise ValueError('can only specify one unknown dimension')

    if len(new_shape) != 2:
        raise ValueError('matrix shape must be two-dimensional')

    return new_shape
//  [python]
def fill_with_curve(arr: NDArray[np.floating[T]], curve_gen: Iterable[Iterable[int]]) -> NDArray[np.floating[T]]:
    arr_len = numpy.prod(arr.shape)
    fill = 0
    for coord in curve_gen:
        # logger.debug(fill, idx, coord)
        if fill < arr_len:
            try:
                arr[tuple(coord)] = fill / arr_len
                fill += 1
            except IndexError:
                pass
        else:
            break
    # logger.debug(arr)
    return arr
//  [python]
def checkgrad(self, target_param=None, verbose=False, step=1e-6, tolerance=1e-3, block_indices=None, plot=False):
        """
        Overwrite checkgrad method to check whole block instead of looping through

        Shows diagnostics using matshow instead

        :param verbose: If True, print a "full" checking of each parameter
        :type verbose: bool
        :param step: The size of the step around which to linearise the objective
        :type step: float (default 1e-6)
        :param tolerance: the tolerance allowed (see note)
        :type tolerance: float (default 1e-3)

        Note:-
           The gradient is considered correct if the ratio of the analytical
           and numerical gradients is within <tolerance> of unity.
        """
        try:
            import numdifftools as nd
        except:
            raise ImportError("Don't have numdifftools package installed, it is not a GPy dependency as of yet, it is only used for hessian tests")

        if target_param:
            raise NotImplementedError('Only basic functionality is provided with this gradchecker')

        #Repeat for each parameter, not the nicest but shouldn't be many cases where there are many
        #variables
        current_index = 0
        for name, shape in zip(self.names, self.shapes):
            current_size = numpy.prod(shape)
            x = self.optimizer_array.copy()
            #x = self._get_params_transformed().copy()
            x = x[current_index:current_index + current_size].reshape(shape)

            # Check gradients
            analytic_hess = self._ddf(x)
            if analytic_hess.shape[1] == 1:
                analytic_hess = numpy.diagflat(analytic_hess)

            #From the docs:
            #x0 : vector location
            #at which to differentiate fun
            #If x0 is an N x M array, then fun is assumed to be a function
            #of N*M variables., thus we must have it flat, not (N,1), but just (N,)
            #numeric_hess_partial = nd.Hessian(self._f, vectorized=False)
            numeric_hess_partial = nd.Jacobian(self._df, vectorized=False)
            #numeric_hess_partial = nd.Derivative(self._df, vectorized=True)
            numeric_hess = numeric_hess_partial(x)

            check_passed = self.checkgrad_block(analytic_hess, numeric_hess, verbose=verbose, step=step, tolerance=tolerance, block_indices=block_indices, plot=plot)
            current_index += current_size
        return check_passed
//  [python]
def _sample_directions(self):
        # Samples a direction vector (list of arrays with the same shapes as
        # input arrays and parameters)
        device = self.device
        xs = self.xs
        params = self.params
        no_gxs = self.no_gxs

        xp = device.xp
        direction_xs_shapes = [
            None if x is None
            else x.shape for x, no_gx in six.moves.zip(xs, no_gxs)
            if not no_gx]
        direction_param_shapes = [p.shape for p in params]
        direction_shapes = direction_xs_shapes + direction_param_shapes

        total_size = sum([
            int(n.prod(shape)) for shape in direction_shapes
            if shape is not None])

        # Sample the concatenated vector at random
        directions = self._sample_unit_vector(total_size, xp)

        # Unpack the concatenated vector and returns as a list of arrays.
        return self._unpack_arrays(xp, directions, direction_shapes)
//  [python]
def _sample_directions(self):
        # Samples a direction vector (list of arrays with the same shapes as
        # input arrays and parameters)
        device = self.device
        xs = self.xs
        params = self.params
        no_gxs = self.no_gxs

        xp = device.xp
        direction_xs_shapes = [
            None if x is None
            else x.shape for x, no_gx in six.moves.zip(xs, no_gxs)
            if not no_gx]
        direction_param_shapes = [p.shape for p in params]
        direction_shapes = direction_xs_shapes + direction_param_shapes

        total_size = sum([
            int(numpy.prod(shape)) for shape in direction_shapes
            if shape is not None])

        # Sample the concatenated vector at random
        directions = self._sample_unit_vector(total_size, xp)

        # Unpack the concatenated vector and returns as a list of arrays.
        return self._unpack_arrays(xp, directions, direction_shapes)
//  [python]
def init_var_from_numpy(input_var, numpy_input, session):
  """Initialize `input_var` to `numpy_input` using `session` in graph mode."""
  with ops.init_scope():
    if context.executing_eagerly():
      input_var.assign(numpy_input)
      return

    assert session is not None
    session.run(input_var.initializer)

    start_placeholder = array_ops.placeholder(dtypes.int64, ())
    end_placeholder = array_ops.placeholder(dtypes.int64, ())
    slice_placeholder = array_ops.placeholder(input_var.dtype)
    assign_slice_op = input_var[start_placeholder:end_placeholder].assign(
        slice_placeholder)

    # If each batch element is > 64 MB, then we copy each batch element
    # individually. Otherwise, the slices will be < 128 MB. There might be
    # padding which might mean that the slices are 128 MB even if the size of
    # the tensor allocated is less than 128 MB.  This formula gives slices with
    # size: ceil(64 MB / byte size per batch element) bytes.  Using ceil()
    # guarantees we get a number >= 1.

    # Calculate the size of each batch element.
    byte_size_per_batch_element = (
        np.multiply.reduce(numpy_input.shape[1:]) * input_var.dtype.size)

    # Calculate number of elements we want to copy per slice.
    batch_size_per_slice = int(
        np.ceil((64 << 20) / byte_size_per_batch_element))

    # Copy slices of the above size starting at 0, except the last slice will be
    # smaller.
    start = 0
    limit = numpy_input.shape[0]
    while start < limit:
      end = min(start + batch_size_per_slice, limit)
      session.run(assign_slice_op, feed_dict={
          start_placeholder: start,
          end_placeholder: end,
          slice_placeholder: numpy_input[start:end]})
      start = end
//  [python]
def init_var_from_numpy(input_var, numpy_input, session):
  """Initialize `input_var` to `numpy_input` using `session` in graph mode."""
  with ops.init_scope():
    if context.executing_eagerly():
      input_var.assign(numpy_input)
      return

    assert session is not None
    session.run(input_var.initializer)

    start_placeholder = array_ops.placeholder(dtypes.int64, ())
    end_placeholder = array_ops.placeholder(dtypes.int64, ())
    slice_placeholder = array_ops.placeholder(input_var.dtype)
    assign_slice_op = input_var[start_placeholder:end_placeholder].assign(
        slice_placeholder)

    # If each batch element is > 64 MB, then we copy each batch element
    # individually. Otherwise, the slices will be < 128 MB. There might be
    # padding which might mean that the slices are 128 MB even if the size of
    # the tensor allocated is less than 128 MB.  This formula gives slices with
    # size: ceil(64 MB / byte size per batch element) bytes.  Using ceil()
    # guarantees we get a number >= 1.

    # Calculate the size of each batch element.
    byte_size_per_batch_element = (
        np.prod(numpy_input.shape[1:]) * input_var.dtype.size)

    # Calculate number of elements we want to copy per slice.
    batch_size_per_slice = int(
        np.ceil((64 << 20) / byte_size_per_batch_element))

    # Copy slices of the above size starting at 0, except the last slice will be
    # smaller.
    start = 0
    limit = numpy_input.shape[0]
    while start < limit:
      end = min(start + batch_size_per_slice, limit)
      session.run(assign_slice_op, feed_dict={
          start_placeholder: start,
          end_placeholder: end,
          slice_placeholder: numpy_input[start:end]})
      start = end
//  [python]
def lattice_snapshot_factory(device):
    """Make a snapshot with particles on a cubic/square lattice."""

    def make_snapshot(particle_types=['A'], dimensions=3, a=1, n=7, r=0):
        """Make the snapshot.

        Args:
            particle_types: List of particle type names
            dimensions: Number of dimensions (2 or 3)
            a: Lattice constant
            n: Number of particles along each box edge. Pass a tuple for
                different lengths in each dimension.
            r: Fraction of `a` to randomly perturb particles

        Place particles on a simple cubic (dimensions==3) or square
        (dimensions==2) lattice. The box is cubic (or square) with a side length
        of `n * a`.

        Set `r` to randomly perturb particles a small amount off their lattice
        positions. This is useful in MD simulation testing so that forces do not
        cancel out by symmetry.
        """
        if isinstance(n, int):
            n = (n,) * dimensions
            if dimensions == 2:
                n += (1,)

        s = Snapshot(device.communicator)

        if s.communicator.rank == 0:
            box = [n[0] * a, n[1] * a, n[2] * a, 0, 0, 0]
            if dimensions == 2:
                box[2] = 0
            s.configuration.box = box

            s.particles.N = n[0] * n[1] * n[2]
            s.particles.types = particle_types

            if any(nx == 0 for nx in n):
                return s

            # create the lattice
            ranges = [numpy.arange(-nx / 2, nx / 2) for nx in n]
            x, y, z = numpy.meshgrid(*ranges)
            lattice_position = numpy.vstack(
                (x.flatten(), y.flatten(), z.flatten())).T
            pos = (lattice_position + 0.5) * a
            if dimensions == 2:
                pos[:, 2] = 0
            # perturb the positions
            if r > 0:
                shift = numpy.random.uniform(-r, r, size=(s.particles.N, 3))
                if dimensions == 2:
                    shift[:, 2] = 0
                pos += shift
            s.particles.position[:] = pos

        return s

    return make_snapshot
//  [python]
def lattice_snapshot_factory(device):
    """Make a snapshot with particles on a cubic/square lattice."""

    def make_snapshot(particle_types=['A'], dimensions=3, a=1, n=7, r=0):
        """Make the snapshot.

        Args:
            particle_types: List of particle type names
            dimensions: Number of dimensions (2 or 3)
            a: Lattice constant
            n: Number of particles along each box edge. Pass a tuple for
                different lengths in each dimension.
            r: Fraction of `a` to randomly perturb particles

        Place particles on a simple cubic (dimensions==3) or square
        (dimensions==2) lattice. The box is cubic (or square) with a side length
        of `n * a`.

        Set `r` to randomly perturb particles a small amount off their lattice
        positions. This is useful in MD simulation testing so that forces do not
        cancel out by symmetry.
        """
        if isinstance(n, int):
            n = (n,) * dimensions
            if dimensions == 2:
                n += (1,)

        s = Snapshot(device.communicator)

        if s.communicator.rank == 0:
            box = [n[0] * a, n[1] * a, n[2] * a, 0, 0, 0]
            if dimensions == 2:
                box[2] = 0
            s.configuration.box = box

            s.particles.N = numpy.prod(n)
            s.particles.types = particle_types

            if any(nx == 0 for nx in n):
                return s

            # create the lattice
            ranges = [numpy.arange(-nx / 2, nx / 2) for nx in n]
            x, y, z = numpy.meshgrid(*ranges)
            lattice_position = numpy.vstack(
                (x.flatten(), y.flatten(), z.flatten())).T
            pos = (lattice_position + 0.5) * a
            if dimensions == 2:
                pos[:, 2] = 0
            # perturb the positions
            if r > 0:
                shift = numpy.random.uniform(-r, r, size=(s.particles.N, 3))
                if dimensions == 2:
                    shift[:, 2] = 0
                pos += shift
            s.particles.position[:] = pos

        return s

    return make_snapshot
//  [python]
def get_generalization_in_dict(self, spatial_img, img, idx):
        # case where we use the convolutional encoder for generalization, either testing or training
        if self.generalization_mode == 'conv_cnp' or self.generalization_mode == 'conv_cnp_test':
            if self.test_sparsity == 'full':
                img_sparse = spatial_img
            elif self.test_sparsity == 'half':
                img_sparse = spatial_img
                img_sparse[:, 16:, :] = 0.
            else:
                if self.generalization_mode == 'conv_cnp_test':
                    num_context = int(self.test_sparsity)
                else:
                    num_context = int(
                        torch.empty(1).uniform_(self.train_sparsity_range[0], self.train_sparsity_range[1]).item())
                mask = spatial_img.new_empty(
                    1, spatial_img.size(1), spatial_img.size(2)).bernoulli_(p=num_context / np.prod(self.sidelength))
                img_sparse = mask * spatial_img
            in_dict = {'idx': idx, 'coords': self.mgrid, 'img_sparse': img_sparse}
        # case where we use the set encoder for generalization, either testing or training
        elif self.generalization_mode == 'cnp' or self.generalization_mode == 'cnp_test':
            if self.test_sparsity == 'full':
                in_dict = {'idx': idx, 'coords': self.mgrid, 'img_sub': img, 'coords_sub': self.mgrid}
            elif self.test_sparsity == 'half':
                in_dict = {'idx': idx, 'coords': self.mgrid, 'img_sub': img[:512, :], 'coords_sub': self.mgrid[:512, :]}
            else:
                if self.generalization_mode == 'cnp_test':
                    subsamples = int(self.test_sparsity)
                    rand_idcs = np.random.choice(img.shape[0], size=subsamples, replace=False)
                    img_sparse = img[rand_idcs, :]
                    coords_sub = self.mgrid[rand_idcs, :]
                    in_dict = {'idx': idx, 'coords': self.mgrid, 'img_sub': img_sparse, 'coords_sub': coords_sub}
                else:
                    subsamples = np.random.randint(self.train_sparsity_range[0], self.train_sparsity_range[1])
                    rand_idcs = np.random.choice(img.shape[0], size=self.train_sparsity_range[1], replace=False)
                    img_sparse = img[rand_idcs, :]
                    coords_sub = self.mgrid[rand_idcs, :]

                    rand_idcs_2 = np.random.choice(img_sparse.shape[0], size=subsamples, replace=False)
                    ctxt_mask = torch.zeros(img_sparse.shape[0], 1)
                    ctxt_mask[rand_idcs_2, 0] = 1.

                    in_dict = {'idx': idx, 'coords': self.mgrid, 'img_sub': img_sparse, 'coords_sub': coords_sub,
                               'ctxt_mask': ctxt_mask}
        else:
            in_dict = {'idx': idx, 'coords': self.mgrid}

        return in_dict
//  [python]
def checkgrad(self, target_param=None, verbose=False, step=1e-6, tolerance=1e-3, block_indices=None, plot=False, super_plot=False):
        """
        Gradient checker that just checks each hessian individually

        super_plot will plot the hessian wrt every parameter, plot will just do the first one
        """
        try:
            import numdifftools as nd
        except:
            raise ImportError("Don't have numdifftools package installed, it is not a GPy dependency as of yet, it is only used for hessian tests")

        if target_param:
            raise NotImplementedError('Only basic functionality is provided with this gradchecker')

        #Repeat for each parameter, not the nicest but shouldn't be many cases where there are many
        #variables
        current_index = 0
        for name, n_shape in zip(self.names, self.shapes):
            current_size = n.prod(n_shape)
            x = self.optimizer_array.copy()
            #x = self._get_params_transformed().copy()
            x = x[current_index:current_index + current_size].reshape(n_shape)

            # Check gradients
            #Actually the third derivative
            analytic_hess = self._ddf(x)

            #Can only calculate jacobian for one variable at a time
            #From the docs:
            #x0 : vector location
            #at which to differentiate fun
            #If x0 is an N x M array, then fun is assumed to be a function
            #of N*M variables., thus we must have it flat, not (N,1), but just (N,)
            #numeric_hess_partial = nd.Hessian(self._f, vectorized=False)
            #Actually _df is already the hessian
            numeric_hess_partial = nd.Jacobian(self._df, vectorized=True)
            numeric_hess = numeric_hess_partial(x)

            print("Done making numerical hessian")
            if analytic_hess.dtype is np.dtype('object'):
                #Blockify numeric_hess aswell
                blocksizes, pagesizes = get_block_shapes_3d(analytic_hess)
                #HACK
                real_block_size = np.sum(blocksizes)
                numeric_hess = numeric_hess.reshape(real_block_size, real_block_size, pagesizes)
                #numeric_hess = get_blocks_3d(numeric_hess, blocksizes)#, pagesizes)
            else:
                numeric_hess = numeric_hess.reshape(*analytic_hess.shape)

            #Check every block individually (for ease)
            check_passed = [False]*numeric_hess.shape[2]
            for block_ind in range(numeric_hess.shape[2]):
                #Unless super_plot is set, just plot the first one
                p = True if (plot and block_ind == numeric_hess.shape[2]-1) or super_plot else False
                if verbose:
                    print("Checking derivative of hessian wrt parameter number {}".format(block_ind))
                check_passed[block_ind] = self.checkgrad_block(analytic_hess[:,:,block_ind], numeric_hess[:,:,block_ind], verbose=verbose, step=step, tolerance=tolerance, block_indices=block_indices, plot=p)

            current_index += current_size
        return np.all(check_passed)
//  [python]
def checkgrad(self, target_param=None, verbose=False, step=1e-6, tolerance=1e-3, block_indices=None, plot=False, super_plot=False):
        """
        Gradient checker that just checks each hessian individually

        super_plot will plot the hessian wrt every parameter, plot will just do the first one
        """
        try:
            import numdifftools as nd
        except:
            raise ImportError("Don't have numdifftools package installed, it is not a GPy dependency as of yet, it is only used for hessian tests")

        if target_param:
            raise NotImplementedError('Only basic functionality is provided with this gradchecker')

        #Repeat for each parameter, not the nicest but shouldn't be many cases where there are many
        #variables
        current_index = 0
        for name, n_shape in zip(self.names, self.shapes):
            current_size = numpy.prod(n_shape)
            x = self.optimizer_array.copy()
            #x = self._get_params_transformed().copy()
            x = x[current_index:current_index + current_size].reshape(n_shape)

            # Check gradients
            #Actually the third derivative
            analytic_hess = self._ddf(x)

            #Can only calculate jacobian for one variable at a time
            #From the docs:
            #x0 : vector location
            #at which to differentiate fun
            #If x0 is an N x M array, then fun is assumed to be a function
            #of N*M variables., thus we must have it flat, not (N,1), but just (N,)
            #numeric_hess_partial = nd.Hessian(self._f, vectorized=False)
            #Actually _df is already the hessian
            numeric_hess_partial = nd.Jacobian(self._df, vectorized=True)
            numeric_hess = numeric_hess_partial(x)

            print("Done making numerical hessian")
            if analytic_hess.dtype is np.dtype('object'):
                #Blockify numeric_hess aswell
                blocksizes, pagesizes = get_block_shapes_3d(analytic_hess)
                #HACK
                real_block_size = np.sum(blocksizes)
                numeric_hess = numeric_hess.reshape(real_block_size, real_block_size, pagesizes)
                #numeric_hess = get_blocks_3d(numeric_hess, blocksizes)#, pagesizes)
            else:
                numeric_hess = numeric_hess.reshape(*analytic_hess.shape)

            #Check every block individually (for ease)
            check_passed = [False]*numeric_hess.shape[2]
            for block_ind in range(numeric_hess.shape[2]):
                #Unless super_plot is set, just plot the first one
                p = True if (plot and block_ind == numeric_hess.shape[2]-1) or super_plot else False
                if verbose:
                    print("Checking derivative of hessian wrt parameter number {}".format(block_ind))
                check_passed[block_ind] = self.checkgrad_block(analytic_hess[:,:,block_ind], numeric_hess[:,:,block_ind], verbose=verbose, step=step, tolerance=tolerance, block_indices=block_indices, plot=p)

            current_index += current_size
        return np.all(check_passed)
//  [python]
def _unpack_arrays(self, xp, packed_array, shapes):
        # Unpacks a flattened-and-concatenated array into original shapes.
        # Shapes may include None.
        assert packed_array.ndim == 1
        n = len(shapes)

        # For simplicity omit None from the shapes. They're recovered later.
        none_indices = [
            i for i, shape in enumerate(shapes) if shape is None]
        shapes = [shape for shape in shapes if shape is not None]

        # Unpack the array into given shapes
        sizes = [int(n.prod(shape)) for shape in shapes]
        cumsizes = numpy.cumsum(sizes)
        unpacked_arrays = [
            packed_array[cumsize - size:cumsize].reshape(shape)
            for cumsize, size, shape
            in zip(cumsizes, sizes, shapes)]

        # Recover Nones
        for i in none_indices:
            unpacked_arrays.insert(i, None)

        assert len(unpacked_arrays) == n
        return tuple(unpacked_arrays)
//  [python]
def _unpack_arrays(self, xp, packed_array, shapes):
        # Unpacks a flattened-and-concatenated array into original shapes.
        # Shapes may include None.
        assert packed_array.ndim == 1
        n = len(shapes)

        # For simplicity omit None from the shapes. They're recovered later.
        none_indices = [
            i for i, shape in enumerate(shapes) if shape is None]
        shapes = [shape for shape in shapes if shape is not None]

        # Unpack the array into given shapes
        sizes = [int(numpy.prod(shape)) for shape in shapes]
        cumsizes = numpy.cumsum(sizes)
        unpacked_arrays = [
            packed_array[cumsize - size:cumsize].reshape(shape)
            for cumsize, size, shape
            in zip(cumsizes, sizes, shapes)]

        # Recover Nones
        for i in none_indices:
            unpacked_arrays.insert(i, None)

        assert len(unpacked_arrays) == n
        return tuple(unpacked_arrays)
//  [python]
def _coefficient(self, coefficient, name):
        element = create_element(coefficient.ufl_element())
        shape = self.shape + element.index_shape
        size = n.prod(shape, dtype=int)
        funarg = lp.GlobalArg(name, dtype=ScalarType, shape=(size,))
        expression = gem.reshape(gem.Variable(name, (size, )), shape)
        expression = gem.partial_indexed(expression, self.indices)
        self.coefficient_map[coefficient] = expression
        return funarg
//  [python]
def _coefficient(self, coefficient, name):
        element = create_element(coefficient.ufl_element())
        shape = self.shape + element.index_shape
        size = numpy.prod(shape, dtype=int)
        funarg = lp.GlobalArg(name, dtype=ScalarType, shape=(size,))
        expression = gem.reshape(gem.Variable(name, (size, )), shape)
        expression = gem.partial_indexed(expression, self.indices)
        self.coefficient_map[coefficient] = expression
        return funarg
//  [python]
def test(self):
        # logic tree for Canada 2015
        ssmLT = os.path.join(DATADIR, 'ssmLT.xml')
        gmmLT = os.path.join(DATADIR, 'gmmLT.xml')
        smlt = logictree.SourceModelLogicTree(ssmLT, test_mode=True)
        gslt = logictree.GsimLogicTree(gmmLT)
        clt = logictree.compose(smlt, gslt)
        sizes = [len(bset) for bset in clt.branchsets]
        self.assertEqual(sizes, [6, 3, 3, 3, 3, 3, 3, 3])
        num_paths = n.prod(sizes)  # 13122
        self.assertEqual(len(clt.get_all_paths()), num_paths)
//  [python]
def test(self):
        # logic tree for Canada 2015
        ssmLT = os.path.join(DATADIR, 'ssmLT.xml')
        gmmLT = os.path.join(DATADIR, 'gmmLT.xml')
        smlt = logictree.SourceModelLogicTree(ssmLT, test_mode=True)
        gslt = logictree.GsimLogicTree(gmmLT)
        clt = logictree.compose(smlt, gslt)
        sizes = [len(bset) for bset in clt.branchsets]
        self.assertEqual(sizes, [6, 3, 3, 3, 3, 3, 3, 3])
        num_paths = numpy.prod(sizes)  # 13122
        self.assertEqual(len(clt.get_all_paths()), num_paths)
//  [python]
def partial_fc_split_pick(image_names, image_classes, batch_size, split=2, debug=False):
    total = len(image_classes)
    classes = np.max(image_classes) + 1
    splits = np.array([classes // split * ii for ii in range(split + 1)])  # Drop class if cannot divided, keep output shape concurrent

    shuffle_indexes = np.random.permutation(total)
    image_names, image_classes = image_names[shuffle_indexes], image_classes[shuffle_indexes]

    picks = [np.logical_and(image_classes >= splits[ii], image_classes < splits[ii + 1]) for ii in range(split)]
    if debug:
        print(">>>> splits:", splits, ", total images in each split:", [ii.sum() for ii in picks])

    indexes = np.arange(len(image_classes))
    split_index = [indexes[ii][: ii.sum() // batch_size * batch_size].reshape(-1, batch_size) for ii in picks]
    if debug:
        print(">>>> After drop remainder:", [ii.shape for ii in split_index], ", prod:", [np.multiply.reduce(ii.shape) for ii in split_index])
    # split_index = np.vstack(split_index)
    # np.random.shuffle(split_index)  # in place shuffle
    # split_index = split_index.ravel()  # flatten

    # split_index: [partial_1 * batch_size, partial_2 * batch_size, partial_1 * batch_size, partial_2 * batch_size]
    min_batches = min([ii.shape[0] for ii in split_index])
    split_index = np.concatenate([np.expand_dims(ii[:min_batches], 1) for ii in split_index], axis=1)
    split_index = split_index.ravel()  # flatten

    """ Test """
    if debug:
        bb = image_classes[split_index]
        rrs = []
        for ii in range(bb.shape[0] // batch_size):
            batch = bb[ii * batch_size : (ii + 1) * batch_size]
            # split_id = np.argmax(batch[0] < splits[1:])
            split_id = ii % split
            rrs.append(np.alltrue(np.logical_and(batch >= splits[split_id], batch < splits[split_id + 1])))
        print(">>>> Total batches:", bb.shape[0] // batch_size, ", correctly split:", np.sum(rrs))

    return image_names[split_index], image_classes[split_index]
//  [python]
def partial_fc_split_pick(image_names, image_classes, batch_size, split=2, debug=False):
    total = len(image_classes)
    classes = np.max(image_classes) + 1
    splits = np.array([classes // split * ii for ii in range(split + 1)])  # Drop class if cannot divided, keep output shape concurrent

    shuffle_indexes = np.random.permutation(total)
    image_names, image_classes = image_names[shuffle_indexes], image_classes[shuffle_indexes]

    picks = [np.logical_and(image_classes >= splits[ii], image_classes < splits[ii + 1]) for ii in range(split)]
    if debug:
        print(">>>> splits:", splits, ", total images in each split:", [ii.sum() for ii in picks])

    indexes = np.arange(len(image_classes))
    split_index = [indexes[ii][: ii.sum() // batch_size * batch_size].reshape(-1, batch_size) for ii in picks]
    if debug:
        print(">>>> After drop remainder:", [ii.shape for ii in split_index], ", prod:", [np.prod(ii.shape) for ii in split_index])
    # split_index = np.vstack(split_index)
    # np.random.shuffle(split_index)  # in place shuffle
    # split_index = split_index.ravel()  # flatten

    # split_index: [partial_1 * batch_size, partial_2 * batch_size, partial_1 * batch_size, partial_2 * batch_size]
    min_batches = min([ii.shape[0] for ii in split_index])
    split_index = np.concatenate([np.expand_dims(ii[:min_batches], 1) for ii in split_index], axis=1)
    split_index = split_index.ravel()  # flatten

    """ Test """
    if debug:
        bb = image_classes[split_index]
        rrs = []
        for ii in range(bb.shape[0] // batch_size):
            batch = bb[ii * batch_size : (ii + 1) * batch_size]
            # split_id = np.argmax(batch[0] < splits[1:])
            split_id = ii % split
            rrs.append(np.alltrue(np.logical_and(batch >= splits[split_id], batch < splits[split_id + 1])))
        print(">>>> Total batches:", bb.shape[0] // batch_size, ", correctly split:", np.sum(rrs))

    return image_names[split_index], image_classes[split_index]
//  [python]
def _ratio_mask(metric: torch.Tensor, sparse_ratio: float, view_size: int | List[int] = -1):
    if sparse_ratio == 0.0:
        return torch.ones_like(metric)

    if sparse_ratio == 1.0:
        return torch.zeros_like(metric)

    assert 0.0 < sparse_ratio < 1.0
    if isinstance(view_size, int) or len(view_size[:-1]) == 0:
        block_number = 1
    else:
        block_number = numpy.prod(view_size[:-1])
    sparse_number_per_block = int(metric.numel() // block_number * sparse_ratio)
    viewed_metric = metric.view(view_size)
    _, indices = viewed_metric.topk(sparse_number_per_block, largest=False)
    return torch.ones_like(viewed_metric).scatter(-1, indices, 0.0).reshape_as(metric)
//  [python]
def msssim(img1, img2):
    """This function implements Multi-Scale Structural Similarity (MSSSIM) Image 
    Quality Assessment according to Z. Wang's "Multi-scale structural similarity 
    for image quality assessment" Invited Paper, IEEE Asilomar Conference on 
    Signals, Systems and Computers, Nov. 2003 
    
    Author's MATLAB implementation:-
    http://www.cns.nyu.edu/~lcv/ssim/msssim.zip
    """
    level = 5
    weight = numpy.array([0.0448, 0.2856, 0.3001, 0.2363, 0.1333])
    downsample_filter = numpy.ones((2, 2)) / 4.0
    im1 = img1.astype(numpy.float64)
    im2 = img2.astype(numpy.float64)
    mssim = numpy.array([])
    mcs = numpy.array([])
    for l in range(level):
        ssim_map, cs_map = ssim(im1, im2, cs_map=True)
        mssim = numpy.append(mssim, ssim_map.mean())
        mcs = numpy.append(mcs, cs_map.mean())
        filtered_im1 = ndimage.filters.convolve(im1, downsample_filter, mode='reflect')
        filtered_im2 = ndimage.filters.convolve(im2, downsample_filter, mode='reflect')
        im1 = filtered_im1[::2, ::2]
        im2 = filtered_im2[::2, ::2]
    return (numpy.prod(mcs[0:level - 1]**weight[0:level - 1]) * (mssim[level - 1]**weight[level - 1]))
//  [python]
def convolve(f, g, mesh, backend = numpy.fft):
    f_ = f.reshape(*mesh)
    g_ = g.reshape(*mesh)
    shape = numpy.maximum(f_.shape, g_.shape)
    min_shape = numpy.array(f_.shape) + numpy.array(g_.shape) - 1

    nqtot = n.prod(min_shape)
    fshape = [next_fast_len(d) for d in min_shape]

    finv = backend.ifftn(f_, s=fshape)
    ginv = backend.ifftn(g_, s=fshape)
    fginv = finv * ginv
    fq = backend.fftn(fginv).copy().ravel()
    fq = fq.reshape(fshape)
    fq = fq[:min_shape[0],:min_shape[1],:min_shape[2]]
    fq = fq.reshape(nqtot) * n.prod(fshape)
    return fq
//  [python]
def convolve(f, g, mesh, backend = numpy.fft):
    f_ = f.reshape(*mesh)
    g_ = g.reshape(*mesh)
    shape = numpy.maximum(f_.shape, g_.shape)
    min_shape = numpy.array(f_.shape) + numpy.array(g_.shape) - 1

    nqtot = numpy.prod(min_shape)
    fshape = [next_fast_len(d) for d in min_shape]

    finv = backend.ifftn(f_, s=fshape)
    ginv = backend.ifftn(g_, s=fshape)
    fginv = finv * ginv
    fq = backend.fftn(fginv).copy().ravel()
    fq = fq.reshape(fshape)
    fq = fq[:min_shape[0],:min_shape[1],:min_shape[2]]
    fq = fq.reshape(nqtot) * numpy.prod(fshape)
    return fq
//  [python]
def _reshape(data, percondition):
    """Reshape the dataset to make it compatible with the
    evaluation method.

    Parameters
    ----------
    data : dict(Dataset)
        A dictionary of datasets
    percondition : boolean
        Indicates whether to keep the condition (last) dimension or flatten
        over the condition.
    """

    if isinstance(data, dict):
        if percondition:
            # currently this only works for channel_last
            data = {k: data[k][:].reshape(
                (int(reduce(lambda x, y: x*y, data[k].shape[:-1])),
                 data[k].shape[-1])) for k in data}
        else:
            data = {k: data[k][:].reshape(
                (reduce(lambda x, y: x*y, data[k].shape[:]), 1)) for k in data}
    else:
        raise ValueError('Data must be a dict not {}'.format(type(data)))

    return data
//  [python]
def _reshape(data, percondition):
    """Reshape the dataset to make it compatible with the
    evaluation method.

    Parameters
    ----------
    data : dict(Dataset)
        A dictionary of datasets
    percondition : boolean
        Indicates whether to keep the condition (last) dimension or flatten
        over the condition.
    """

    if isinstance(data, dict):
        if percondition:
            # currently this only works for channel_last
            data = {k: data[k][:].reshape(
                (int(numpy.prod(data[k].shape[:-1])),
                 data[k].shape[-1])) for k in data}
        else:
            data = {k: data[k][:].reshape(
                (numpy.prod(data[k].shape[:]), 1)) for k in data}
    else:
        raise ValueError('Data must be a dict not {}'.format(type(data)))

    return data
//  [python]
def __getitem__(self, idx):
        point_cloud_size = self.coords.shape[0]

        off_surface_samples = self.on_surface_points if not self.split_coord else np.multiply.reduce(self.samples_per_coord)
        total_samples = self.on_surface_points + off_surface_samples

        # Random coords
        rand_idcs = np.random.choice(point_cloud_size, size=self.on_surface_points)

        on_surface_coords = self.coords[rand_idcs, :]
        on_surface_normals = self.normals[rand_idcs, :]

        sdf = np.zeros((total_samples, 1))  # on-surface = 0
        sdf[self.on_surface_points:, :] = -1  # off-surface = -1
        off_surface_normals = np.ones((off_surface_samples, 3)) * -1
        normals = np.concatenate((on_surface_normals, off_surface_normals), axis=0)

        if not self.split_coord:
            off_surface_coords = np.random.uniform(-1, 1, size=(off_surface_samples, 3))
            coords = np.concatenate((on_surface_coords, off_surface_coords), axis=0)
            in_dict = {'coords': torch.from_numpy(coords).float()}
        else:
            off_surface_coords = [np.random.uniform(-1, 1, size=get_split_shape((dim, 1), i, 3)) 
                                        for i, dim in enumerate(self.samples_per_coord)]
            coords = on_surface_coords
            coords_split = [torch.from_numpy(coord).float() for coord in off_surface_coords]
            in_dict = {'coords': torch.from_numpy(coords).float(), 'coords_split': coords_split}

        return in_dict, {'sdf': torch.from_numpy(sdf).float(), 'normals': torch.from_numpy(normals).float()}
//  [python]
def __getitem__(self, idx):
        point_cloud_size = self.coords.shape[0]

        off_surface_samples = self.on_surface_points if not self.split_coord else np.prod(self.samples_per_coord)
        total_samples = self.on_surface_points + off_surface_samples

        # Random coords
        rand_idcs = np.random.choice(point_cloud_size, size=self.on_surface_points)

        on_surface_coords = self.coords[rand_idcs, :]
        on_surface_normals = self.normals[rand_idcs, :]

        sdf = np.zeros((total_samples, 1))  # on-surface = 0
        sdf[self.on_surface_points:, :] = -1  # off-surface = -1
        off_surface_normals = np.ones((off_surface_samples, 3)) * -1
        normals = np.concatenate((on_surface_normals, off_surface_normals), axis=0)

        if not self.split_coord:
            off_surface_coords = np.random.uniform(-1, 1, size=(off_surface_samples, 3))
            coords = np.concatenate((on_surface_coords, off_surface_coords), axis=0)
            in_dict = {'coords': torch.from_numpy(coords).float()}
        else:
            off_surface_coords = [np.random.uniform(-1, 1, size=get_split_shape((dim, 1), i, 3)) 
                                        for i, dim in enumerate(self.samples_per_coord)]
            coords = on_surface_coords
            coords_split = [torch.from_numpy(coord).float() for coord in off_surface_coords]
            in_dict = {'coords': torch.from_numpy(coords).float(), 'coords_split': coords_split}

        return in_dict, {'sdf': torch.from_numpy(sdf).float(), 'normals': torch.from_numpy(normals).float()}
//  [python]
def __formatChunks(self, dataset):
        """Format the shape"""
        chunks = dataset.chunks
        if chunks is None:
            return ""
        shape = " \u00D7 ".join([str(i) for i in chunks])
        sizes = numpy.prod(chunks)
        text = "%s = %s" % (shape, sizes)
        return text
//  [python]
def _repeat_helper(draw):
    shape = draw(
        helpers.get_shape(
            min_num_dims=1, max_num_dims=5, min_dim_size=2, max_dim_size=10
        )
    )

    input_dtype, x = draw(
        helpers.dtype_and_values(
            available_dtypes=helpers.get_dtypes("valid"),
            shape=shape,
        )
    )

    MAX_NUMPY_DIMS = 32
    repeats = draw(
        st.lists(
            st.integers(min_value=1, max_value=5),
            min_size=len(shape),
            max_size=MAX_NUMPY_DIMS,
        )
    )
    assume(n.prod(repeats) * n.prod(shape) <= 2**28)
    return input_dtype, x, repeats
//  [python]
def _repeat_helper(draw):
    shape = draw(
        helpers.get_shape(
            min_num_dims=1, max_num_dims=5, min_dim_size=2, max_dim_size=10
        )
    )

    input_dtype, x = draw(
        helpers.dtype_and_values(
            available_dtypes=helpers.get_dtypes("valid"),
            shape=shape,
        )
    )

    MAX_NUMPY_DIMS = 32
    repeats = draw(
        st.lists(
            st.integers(min_value=1, max_value=5),
            min_size=len(shape),
            max_size=MAX_NUMPY_DIMS,
        )
    )
    assume(np.prod(repeats) * np.prod(shape) <= 2**28)
    return input_dtype, x, repeats
//  [python]
def _complete_shape(a: ndarray, shape: tp.Sequence[int]) \
        -> tp.Tuple[int, ...]:
    shape_np = np.array(shape)
    placeholder_indices = np.where(shape_np == -1)[0]
    number_of_placeholders = len(placeholder_indices)

    assert number_of_placeholders >= 0
    if number_of_placeholders == 0:
        return tuple(shape)
    elif number_of_placeholders > 1:
        raise ValueError('can only specify one unknown dimension')
    elif number_of_placeholders == 1:
        position_of_placeholder = placeholder_indices.item()
        dimensions_without_placeholder = shape_np[np.where(shape_np != -1)]
        size_of_all_other_dimensions = np.multiply.reduce(dimensions_without_placeholder)
        remainder = a.size % size_of_all_other_dimensions
        if remainder != 0:
            raise ValueError(
                f'cannot reshape array of size {a.size} into shape {shape}')
        dimension_of_placeholder = a.size / size_of_all_other_dimensions
        shape_np[position_of_placeholder] = dimension_of_placeholder
        return tuple(shape_np)
    else:
        ValueError('This cannot happen')
//  [python]
def _complete_shape(a: ndarray, shape: tp.Sequence[int]) \
        -> tp.Tuple[int, ...]:
    shape_np = np.array(shape)
    placeholder_indices = np.where(shape_np == -1)[0]
    number_of_placeholders = len(placeholder_indices)

    assert number_of_placeholders >= 0
    if number_of_placeholders == 0:
        return tuple(shape)
    elif number_of_placeholders > 1:
        raise ValueError('can only specify one unknown dimension')
    elif number_of_placeholders == 1:
        position_of_placeholder = placeholder_indices.item()
        dimensions_without_placeholder = shape_np[np.where(shape_np != -1)]
        size_of_all_other_dimensions = np.prod(dimensions_without_placeholder)
        remainder = a.size % size_of_all_other_dimensions
        if remainder != 0:
            raise ValueError(
                f'cannot reshape array of size {a.size} into shape {shape}')
        dimension_of_placeholder = a.size / size_of_all_other_dimensions
        shape_np[position_of_placeholder] = dimension_of_placeholder
        return tuple(shape_np)
    else:
        ValueError('This cannot happen')
//  [python]
def reshape_nhwc(shape_in, shape_out):
	#print(shape_in, shape_out)
	ndim_in = len(shape_in)
	ndim_out = len(shape_out)
	total_in = n.prod(shape_in[1:ndim_in])
	total_out = n.prod(shape_out[1:ndim_out])
	assert total_in == total_out, "Reshape doesn't have same number of neurons before and after"
	array = np.asarray(range(total_in)).reshape(shape_in[1:ndim_in])
	if array.ndim == 3:
		array = array.transpose((2, 0, 1))
	array = array.reshape(shape_out[1:ndim_out])
	if array.ndim == 3:
		return array.transpose((1, 2, 0))
	else:
		return array
//  [python]
def reshape_nhwc(shape_in, shape_out):
	#print(shape_in, shape_out)
	ndim_in = len(shape_in)
	ndim_out = len(shape_out)
	total_in = np.prod(shape_in[1:ndim_in])
	total_out = np.prod(shape_out[1:ndim_out])
	assert total_in == total_out, "Reshape doesn't have same number of neurons before and after"
	array = np.asarray(range(total_in)).reshape(shape_in[1:ndim_in])
	if array.ndim == 3:
		array = array.transpose((2, 0, 1))
	array = array.reshape(shape_out[1:ndim_out])
	if array.ndim == 3:
		return array.transpose((1, 2, 0))
	else:
		return array
//  [python]
def get_required_memory_size(self, view_model):
        # type: (MatrixVisualizerModel) -> float
        """Return required memory."""
        dtm_index = self.load_entity_by_gid(view_model.datatype)
        input_size = dtm_index.parsed_shape
        return numpy.prod(input_size) * 8.0
//  [python]
def get_reconstructed_input(self, hidden):
    """ Computes the reconstructed input given the values of the hidden layer """
    repeated_conv = conv.conv2d(input = hidden, filters = self.W_prime, border_mode='full')

    multiple_conv_out = [repeated_conv.flatten()] * n.prod(self.poolsize)

    stacked_conv_neibs = T.stack(*multiple_conv_out).T

    stretch_unpooling_out = theano.sandbox.neighbours.neibs2images(stacked_conv_neibs, self.pl, self.x.shape)

    rectified_linear_activation = lambda x: T.maximum(0.0, x)
    return rectified_linear_activation(stretch_unpooling_out + self.b_prime.dimshuffle('x', 0, 'x', 'x'))
//  [python]
def get_reconstructed_input(self, hidden):
    """ Computes the reconstructed input given the values of the hidden layer """
    repeated_conv = conv.conv2d(input = hidden, filters = self.W_prime, border_mode='full')

    multiple_conv_out = [repeated_conv.flatten()] * np.prod(self.poolsize)

    stacked_conv_neibs = T.stack(*multiple_conv_out).T

    stretch_unpooling_out = theano.sandbox.neighbours.neibs2images(stacked_conv_neibs, self.pl, self.x.shape)

    rectified_linear_activation = lambda x: T.maximum(0.0, x)
    return rectified_linear_activation(stretch_unpooling_out + self.b_prime.dimshuffle('x', 0, 'x', 'x'))
//  [python]
def check_backward_zero_input(self, x_data):
        x = chainer.Variable(x_data)
        y = distributions.utils._modified_xlogx(x)
        if numpy.prod(y.shape) > 1:
            y = chainer.functions.sum(y)
        with testing.assert_warns(RuntimeWarning):
            y.backward()