jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _get_nan_decs(self):
        """Set all variables to nan for specializations of fused types for
        which don't have signatures.

        """
        # Set non fused-type variables to nan
        tab = " "*4
        fused_types, lines = [], [tab + "else:"]
        seen = set()
        for outvar, outtype, code in zip(self.outvars, self.outtypes,
                                         self.outcodes):
            if len(code) == 1:
                line = "{}[0] = {}".format(outvar, NAN_VALUE[code])
                lines.append(2*tab + line)
            else:
                fused_type = outtype
                name, _ = fused_type
                if name not in seen:
                    fused_types.append(fused_type)
                    seen.add(name)
        if not fused_types:
            return lines

        # Set fused-type variables to nan
        all_codes = tuple([codes for _unused, codes in fused_types])

        codelens = [len(x) for x in all_codes]
        last = reduce(lambda x, y: x*y, codelens) - 1
        for m, codes in enumerate(itertools.product(*all_codes)):
            fused_codes, decs = [], []
            for n, fused_type in enumerate(fused_types):
                code = codes[n]
                fused_codes.append(underscore(CY_TYPES[code]))
                for nn, outvar in enumerate(self.outvars):
                    if self.outtypes[nn] == fused_type:
                        line = "{}[0] = {}".format(outvar, NAN_VALUE[code])
                        decs.append(line)
            if m == 0:
                adverb = "if"
            elif m == last:
                adverb = "else"
            else:
                adverb = "elif"
            cond = self._get_conditional(fused_types, codes, adverb)
            lines.append(2*tab + cond)
            lines.extend([3*tab + x for x in decs])
        return lines
//  [python]
def _get_nan_decs(self):
        """Set all variables to nan for specializations of fused types for
        which don't have signatures.

        """
        # Set non fused-type variables to nan
        tab = " "*4
        fused_types, lines = [], [tab + "else:"]
        seen = set()
        for outvar, outtype, code in zip(self.outvars, self.outtypes,
                                         self.outcodes):
            if len(code) == 1:
                line = "{}[0] = {}".format(outvar, NAN_VALUE[code])
                lines.append(2*tab + line)
            else:
                fused_type = outtype
                name, _ = fused_type
                if name not in seen:
                    fused_types.append(fused_type)
                    seen.add(name)
        if not fused_types:
            return lines

        # Set fused-type variables to nan
        all_codes = tuple([codes for _unused, codes in fused_types])

        codelens = [len(x) for x in all_codes]
        last = numpy.prod(codelens) - 1
        for m, codes in enumerate(itertools.product(*all_codes)):
            fused_codes, decs = [], []
            for n, fused_type in enumerate(fused_types):
                code = codes[n]
                fused_codes.append(underscore(CY_TYPES[code]))
                for nn, outvar in enumerate(self.outvars):
                    if self.outtypes[nn] == fused_type:
                        line = "{}[0] = {}".format(outvar, NAN_VALUE[code])
                        decs.append(line)
            if m == 0:
                adverb = "if"
            elif m == last:
                adverb = "else"
            else:
                adverb = "elif"
            cond = self._get_conditional(fused_types, codes, adverb)
            lines.append(2*tab + cond)
            lines.extend([3*tab + x for x in decs])
        return lines
//  [python]
def result_size(self, input_shape, segment_length, sample_period):
        """
        Returns the storage size in Bytes of the main result (complex array) of 
        the FFT.
        """
        result_size = numpy.prod(self.result_shape(input_shape, segment_length,
                                                   sample_period)) * 2.0 * 8.0  # complex*Bytes
        return result_size
//  [python]
def get_required_memory_size(self, view_model):
        # type: (ICAAdapterModel) -> int
        """
        Return the required memory to run this algorithm.
        """
        used_shape = (self.input_shape[0], 1, self.input_shape[2], self.input_shape[3])
        input_size = numpy.prod(used_shape) * 8.0
        output_size = self.result_size(self.input_shape, view_model.n_components)
        return input_size + output_size
//  [python]
def result_size(self, input_shape, n_components):
        """Returns the storage size in bytes of the mixing matrix of the ICA analysis, assuming 64-bit float."""
        return numpy.prod(self.result_shape(input_shape, n_components)) * 8
//  [python]
def cross_index(values, index):
    """
    Allows efficiently indexing into a cartesian product without
    expanding it. The values should be defined as a list of iterables
    making up the cartesian product and a linear index, returning
    the cross product of the values at the supplied index.
    """
    lengths = [len(v) for v in values]
    length = reduce(lambda x, y: x*y, lengths)
    if index >= length:
        raise IndexError('Index %d out of bounds for cross-product of size %d'
                         % (index, length))
    indexes = []
    for i in range(1, len(values))[::-1]:
        p = reduce(lambda x, y: x*y, lengths[-i:])
        indexes.append(index//p)
        index -= indexes[-1] * p
    indexes.append(index)
    return tuple(v[i] for v, i in zip(values, indexes))
//  [python]
def cross_index(values, index):
    """
    Allows efficiently indexing into a cartesian product without
    expanding it. The values should be defined as a list of iterables
    making up the cartesian product and a linear index, returning
    the cross product of the values at the supplied index.
    """
    lengths = [len(v) for v in values]
    length = np.prod(lengths)
    if index >= length:
        raise IndexError('Index %d out of bounds for cross-product of size %d'
                         % (index, length))
    indexes = []
    for i in range(1, len(values))[::-1]:
        p = np.prod(lengths[-i:])
        indexes.append(index//p)
        index -= indexes[-1] * p
    indexes.append(index)
    return tuple(v[i] for v, i in zip(values, indexes))
//  [python]
def test_reduce_prod(self):
        self.assertEqualVar(expr.reduce_prod(self.x), n.prod(self.x_))
//  [python]
def test_reduce_prod(self):
        self.assertEqualVar(expr.reduce_prod(self.x), np.prod(self.x_))
//  [python]
def _convert_p(self, column, field, recformat):
        """Convert a raw table column of FITS P or Q format descriptors
        to a VLA column with the array data returned from the heap.
        """
        if column.dim:
            vla_shape = tuple(
                reversed(tuple(map(int, column.dim.strip("()").split(","))))
            )
        dummy = _VLF([None] * len(self), dtype=recformat.dtype)
        raw_data = self._get_raw_data()

        if raw_data is None:
            raise OSError(
                f"Could not find heap data for the {column.name!r} variable-length "
                "array column."
            )

        for idx in range(len(self)):
            offset = field[idx, 1] + self._heapoffset
            count = field[idx, 0]

            if recformat.dtype == "S":
                dt = np.dtype(recformat.dtype + str(1))
                arr_len = count * dt.itemsize
                da = raw_data[offset : offset + arr_len].view(dt)
                da = np.char.array(da.view(dtype=dt), itemsize=count)
                dummy[idx] = decode_ascii(da)
            else:
                dt = np.dtype(recformat.dtype)
                arr_len = count * dt.itemsize
                dummy[idx] = raw_data[offset : offset + arr_len].view(dt)
                if column.dim and len(vla_shape) > 1:
                    # The VLA is reshaped consistently with TDIM instructions
                    if vla_shape[0] == 1:
                        dummy[idx] = dummy[idx].reshape(1, len(dummy[idx]))
                    else:
                        vla_dim = vla_shape[1:]
                        vla_first = int(len(dummy[idx]) / reduce(lambda x, y: x*y, vla_dim))
                        dummy[idx] = dummy[idx].reshape((vla_first,) + vla_dim)

                dummy[idx].dtype = dummy[idx].dtype.newbyteorder(">")
                # Each array in the field may now require additional
                # scaling depending on the other scaling parameters
                # TODO: The same scaling parameters apply to every
                # array in the column so this is currently very slow; we
                # really only need to check once whether any scaling will
                # be necessary and skip this step if not
                # TODO: Test that this works for X format; I don't think
                # that it does--the recformat variable only applies to the P
                # format not the X format
                dummy[idx] = self._convert_other(column, dummy[idx], recformat)

        return dummy
//  [python]
def _convert_p(self, column, field, recformat):
        """Convert a raw table column of FITS P or Q format descriptors
        to a VLA column with the array data returned from the heap.
        """
        if column.dim:
            vla_shape = tuple(
                reversed(tuple(map(int, column.dim.strip("()").split(","))))
            )
        dummy = _VLF([None] * len(self), dtype=recformat.dtype)
        raw_data = self._get_raw_data()

        if raw_data is None:
            raise OSError(
                f"Could not find heap data for the {column.name!r} variable-length "
                "array column."
            )

        for idx in range(len(self)):
            offset = field[idx, 1] + self._heapoffset
            count = field[idx, 0]

            if recformat.dtype == "S":
                dt = np.dtype(recformat.dtype + str(1))
                arr_len = count * dt.itemsize
                da = raw_data[offset : offset + arr_len].view(dt)
                da = np.char.array(da.view(dtype=dt), itemsize=count)
                dummy[idx] = decode_ascii(da)
            else:
                dt = np.dtype(recformat.dtype)
                arr_len = count * dt.itemsize
                dummy[idx] = raw_data[offset : offset + arr_len].view(dt)
                if column.dim and len(vla_shape) > 1:
                    # The VLA is reshaped consistently with TDIM instructions
                    if vla_shape[0] == 1:
                        dummy[idx] = dummy[idx].reshape(1, len(dummy[idx]))
                    else:
                        vla_dim = vla_shape[1:]
                        vla_first = int(len(dummy[idx]) / np.prod(vla_dim))
                        dummy[idx] = dummy[idx].reshape((vla_first,) + vla_dim)

                dummy[idx].dtype = dummy[idx].dtype.newbyteorder(">")
                # Each array in the field may now require additional
                # scaling depending on the other scaling parameters
                # TODO: The same scaling parameters apply to every
                # array in the column so this is currently very slow; we
                # really only need to check once whether any scaling will
                # be necessary and skip this step if not
                # TODO: Test that this works for X format; I don't think
                # that it does--the recformat variable only applies to the P
                # format not the X format
                dummy[idx] = self._convert_other(column, dummy[idx], recformat)

        return dummy
//  [python]
def test_advinc_subtensor1():
    """ Test the second case in the opt local_gpu_advanced_incsubtensor1 """
    for shp in [(3, 3), (3, 3, 3)]:
        shared = gpuarray_shared_constructor
        xval = numpy.arange(n.prod(shp), dtype='float32').reshape(shp) + 1
        yval = numpy.empty((2,) + shp[1:], dtype='float32')
        yval[:] = 10
        x = shared(xval, name='x')
        y = tensor.tensor(dtype='float32',
                          broadcastable=(False,) * len(shp),
                          name='y')
        expr = tensor.advanced_inc_subtensor1(x, y, [0, 2])
        f = theano.function([y], expr, mode=mode_with_gpu)
        assert sum([isinstance(node.op, GpuAdvancedIncSubtensor1)
                    for node in f.maker.fgraph.toposort()]) == 1
        rval = f(yval)
        rep = xval.copy()
        rep[[0, 2]] += yval
        assert numpy.allclose(rval, rep)
//  [python]
def test_advinc_subtensor1():
    """ Test the second case in the opt local_gpu_advanced_incsubtensor1 """
    for shp in [(3, 3), (3, 3, 3)]:
        shared = gpuarray_shared_constructor
        xval = numpy.arange(numpy.prod(shp), dtype='float32').reshape(shp) + 1
        yval = numpy.empty((2,) + shp[1:], dtype='float32')
        yval[:] = 10
        x = shared(xval, name='x')
        y = tensor.tensor(dtype='float32',
                          broadcastable=(False,) * len(shp),
                          name='y')
        expr = tensor.advanced_inc_subtensor1(x, y, [0, 2])
        f = theano.function([y], expr, mode=mode_with_gpu)
        assert sum([isinstance(node.op, GpuAdvancedIncSubtensor1)
                    for node in f.maker.fgraph.toposort()]) == 1
        rval = f(yval)
        rep = xval.copy()
        rep[[0, 2]] += yval
        assert numpy.allclose(rval, rep)
//  [python]
def compute_psi1_numpy(lls, lsf, xmean, xvar, z):

    if xmean.ndim == 1:
        xmean = xmean[ None, : ]

    ls = np.exp(lls)
    sf = np.exp(lsf)
    lspxvar = ls + xvar
    constterm1 = ls / lspxvar
    constterm2 = np.prod(np.sqrt(constterm1), 1)
    r2_psi1 = np.outer(np.sum(xmean * xmean / lspxvar, 1), \
        np.ones(z.shape[ 0 ])) - 2 * np.dot(xmean / lspxvar, z.T) + \
        np.dot(1.0 / lspxvar, z.T **2)
    psi1 = sf * np.outer(constterm2, np.ones(z.shape[ 0 ])) * np.exp(-0.5 * r2_psi1)
    return psi1
//  [python]
def process_batch(self, batch):
        """
        Execution of an update step, infer cg_id from selectors, and pick
        corresponding computational graph, and apply batch to the CG.
        """
        cg_id = self.get_cg_id_from_selectors(batch['src_selector'][0],
                                              batch['trg_selector'][0])

        # Apply input replacement with <UNK> if necessary
        if self.drop_input[cg_id] > 0.0:
            num_els = numpy.prod(batch['source'].shape)
            num_reps = max(1, int(num_els * self.drop_input[cg_id]))
            replace_idx = numpy.random.choice(num_els, num_reps, replace=False)
            # TODO: set it according to unk_id in config
            batch['source'][numpy.unravel_index(
                replace_idx, batch['source'].shape)] = 1

        ordered_batch = [batch[v.name] for v in self.algorithms[cg_id].inputs]

        # To save memory, we may combine f_update and f_grad_shared
        if self.f_grad_shareds[cg_id] is None:
            inps = [self.learning_rate] + ordered_batch
            cost = self.f_updates[cg_id](*inps)
            self._cost = ('cost_' + cg_id, cost)
        else:
            cost = self.f_grad_shareds[cg_id](*ordered_batch)
            self._cost = ('cost_' + cg_id, cost)
            self.f_updates[cg_id](self.learning_rate)
//  [python]
def add(self, data, index=None):
        '''
        Adds a batch of data to be considered for the running top k.
        The zeroth dimension enumerates the observations.  All other
        dimensions enumerate different features.
        '''
        if self.top_data is None:
            # Allocation: allocate a buffer of size 5*k, at least 10, for each.
            self.data_shape = data.shape[1:]
            feature_size = int(numpy.prod(self.data_shape))
            self.top_data = torch.zeros(
                feature_size, max(10, self.k * 5), out=data.new())
            self.top_index = self.top_data.clone().long()
            self.linear_index = 0 if len(data.shape) == 1 else torch.arange(
                feature_size, out=self.top_index.new()).mul_(
                self.top_data.shape[-1])[:, None]
        size = data.shape[0]
        sk = min(size, self.k)
        if self.top_data.shape[-1] < self.next + sk:
            # Compression: if full, keep topk only.
            self.top_data[:, :self.k], self.top_index[:, :self.k] = (
                self.result(sorted=False, flat=True))
            self.next = self.k
            free = self.top_data.shape[-1] - self.next
        # Pick: copy the top sk of the next batch into the buffer.
        # Currently strided topk is slow.  So we clone after transpose.
        # TODO: remove the clone() if it becomes faster.
        cdata = data.reshape(size, -1).t().clone()
        td, ti = cdata.topk(sk, sorted=False, largest=self.largest)
        self.top_data[:, self.next:self.next + sk] = td
        if index is not None:
            ti = index[ti]
        else:
            ti = ti + self.count
        self.top_index[:, self.next:self.next + sk] = ti
        self.next += sk
        self.count += size
//  [python]
def count_vars(module: torch.nn.Module) -> int:
    return sum(np.prod(p.shape) for p in module.parameters())
//  [python]
def cartesian(nodes, order="C"):
    """Cartesian product of a list of arrays

    Parameters:
    -----------
    nodes: (list of 1d-arrays)
    order: ('C' or 'F') order in which the product is enumerated

    Returns:
    --------
    out: (2d-array) each line corresponds to one point of the product space
    """

    nodes = [numpy.array(e) for e in nodes]
    shapes = [e.shape[0] for e in nodes]

    n = len(nodes)
    l = reduce(lambda x,y: x*y, shapes)
    out = numpy.zeros((l, n))

    if order == "C":
        repetitions = numpy.cumprod([1] + shapes[:-1])
    else:
        shapes.reverse()
        sh = [1] + shapes[:-1]
        repetitions = numpy.cumprod(sh)
        repetitions = repetitions.tolist()
        repetitions.reverse()

    for i in range(n):
        _repeat_1d(nodes[i], repetitions[i], out[:, i])

    return out
//  [python]
def cartesian(nodes, order="C"):
    """Cartesian product of a list of arrays

    Parameters:
    -----------
    nodes: (list of 1d-arrays)
    order: ('C' or 'F') order in which the product is enumerated

    Returns:
    --------
    out: (2d-array) each line corresponds to one point of the product space
    """

    nodes = [numpy.array(e) for e in nodes]
    shapes = [e.shape[0] for e in nodes]

    n = len(nodes)
    l = numpy.prod(shapes)
    out = numpy.zeros((l, n))

    if order == "C":
        repetitions = numpy.cumprod([1] + shapes[:-1])
    else:
        shapes.reverse()
        sh = [1] + shapes[:-1]
        repetitions = numpy.cumprod(sh)
        repetitions = repetitions.tolist()
        repetitions.reverse()

    for i in range(n):
        _repeat_1d(nodes[i], repetitions[i], out[:, i])

    return out
//  [python]
def k2y(self, k):
        k = self.tf.reshape(k, (numpy.prod(self.Kd), 1))
        # torch_y = self.torch_sp.mv(self.torch.flatten(k))
        tf_y = self.tf.sparse.sparse_dense_matmul(self.tf_sp, k)
        return tf_y[:,0]
//  [python]
def imread(filename):
    """
    Read in an image file with tif or image file type supported by cv2.

    Args:
        filename (str): The path to the image file.

    Returns:
        numpy.ndarray: The image data as a NumPy array.

    Raises:
        None

    Raises an error if the image file format is not supported.

    Examples:
        >>> img = imread("image.tif")
    """
    # ensure that extension check is not case sensitive
    ext = os.path.splitext(filename)[-1].lower()
    if ext == ".tif" or ext == ".tiff":
        with tifffile.TiffFile(filename) as tif:
            ltif = len(tif.pages)
            try:
                full_shape = tif.shaped_metadata[0]["shape"]
            except:
                try:
                    page = tif.series[0][0]
                    full_shape = tif.series[0].shape
                except:
                    ltif = 0
            if ltif < 10:
                img = tif.asarray()
            else:
                page = tif.series[0][0]
                shape, dtype = page.shape, page.dtype
                ltif = int(np.prod(full_shape) / np.prod(shape))
                io_logger.info(f"reading tiff with {ltif} planes")
                img = np.zeros((ltif, *shape), dtype=dtype)
                for i, page in enumerate(tqdm(tif.series[0])):
                    img[i] = page.asarray()
                img = img.reshape(full_shape)
        return img
    elif ext == ".dax":
        img = load_dax(filename)
        return img
    elif ext == ".nd2":
        if not ND2:
            io_logger.critical("ERROR: need to 'pip install nd2' to load in .nd2 file")
            return None
    elif ext == ".nrrd":
        if not NRRD:
            io_logger.critical(
                "ERROR: need to 'pip install pynrrd' to load in .nrrd file")
            return None
        else:
            img, metadata = nrrd.read(filename)
            if img.ndim == 3:
                img = img.transpose(2, 0, 1)
            return img
    elif ext != ".npy":
        try:
            img = cv2.imread(filename, -1)  #cv2.LOAD_IMAGE_ANYDEPTH)
            if img.ndim > 2:
                img = img[..., [2, 1, 0]]
            return img
        except Exception as e:
            io_logger.critical("ERROR: could not read file, %s" % e)
            return None
    else:
        try:
            dat = np.load(filename, allow_pickle=True).item()
            masks = dat["masks"]
            return masks
        except Exception as e:
            io_logger.critical("ERROR: could not read masks from file, %s" % e)
            return None
//  [python]
def get_flat(self):
        return UTPM(self.data.reshape(self.data.shape[:2] + (numpy.prod(self.data.shape[2:]),) ))
//  [python]
def compute_psi2_numpy(lls, lsf, z, input_means, input_vars):

    ls = np.exp(lls)
    sf = np.exp(lsf)
    b = ls / casting(2.0)
    term_1 = np.multiply.reduce(np.sqrt(b / (b + input_vars)), 1)

    scale = np.sqrt(4 * (2 * b[ None, : ] + 0 * input_vars))
    scaled_z = z[ None, : , : ] / scale[ : , None , : ]
    scaled_z_minus_m = scaled_z
    r2b = np.sum(scaled_z_minus_m**2, 2)[ :, None, : ] + np.sum(scaled_z_minus_m**2, 2)[ :, : , None ] - \
        2 * np.einsum('ijk,ikl->ijl', scaled_z_minus_m, np.transpose(scaled_z_minus_m, [ 0, 2, 1 ]))
    term_2 = np.exp(-r2b)

    scale = np.sqrt(4 * (2 * b[ None, : ] + 2 * input_vars))
    scaled_z = z[ None, : , : ] / scale[ : , None , : ]
    scaled_m = input_means / scale
    scaled_m = np.tile(scaled_m[ : , None, : ], [ 1, z.shape[ 0 ], 1])
    scaled_z_minus_m = scaled_z - scaled_m
    r2b = np.sum(scaled_z_minus_m**2, 2)[ :, None, : ] + np.sum(scaled_z_minus_m**2, 2)[ :, : , None ] + \
        2 * np.einsum('ijk,ikl->ijl', scaled_z_minus_m, np.transpose(scaled_z_minus_m, [ 0, 2, 1 ]))
    term_3 = np.exp(-r2b)
    
    psi2_computed = sf**casting(2.0) * term_1[ :, None, None ] * term_2 * term_3
    psi2_computed = np.transpose(psi2_computed, [ 1, 2, 0 ])

    return psi2_computed
//  [python]
def compute_psi2_numpy(lls, lsf, z, input_means, input_vars):

    ls = np.exp(lls)
    sf = np.exp(lsf)
    b = ls / casting(2.0)
    term_1 = np.prod(np.sqrt(b / (b + input_vars)), 1)

    scale = np.sqrt(4 * (2 * b[ None, : ] + 0 * input_vars))
    scaled_z = z[ None, : , : ] / scale[ : , None , : ]
    scaled_z_minus_m = scaled_z
    r2b = np.sum(scaled_z_minus_m**2, 2)[ :, None, : ] + np.sum(scaled_z_minus_m**2, 2)[ :, : , None ] - \
        2 * np.einsum('ijk,ikl->ijl', scaled_z_minus_m, np.transpose(scaled_z_minus_m, [ 0, 2, 1 ]))
    term_2 = np.exp(-r2b)

    scale = np.sqrt(4 * (2 * b[ None, : ] + 2 * input_vars))
    scaled_z = z[ None, : , : ] / scale[ : , None , : ]
    scaled_m = input_means / scale
    scaled_m = np.tile(scaled_m[ : , None, : ], [ 1, z.shape[ 0 ], 1])
    scaled_z_minus_m = scaled_z - scaled_m
    r2b = np.sum(scaled_z_minus_m**2, 2)[ :, None, : ] + np.sum(scaled_z_minus_m**2, 2)[ :, : , None ] + \
        2 * np.einsum('ijk,ikl->ijl', scaled_z_minus_m, np.transpose(scaled_z_minus_m, [ 0, 2, 1 ]))
    term_3 = np.exp(-r2b)
    
    psi2_computed = sf**casting(2.0) * term_1[ :, None, None ] * term_2 * term_3
    psi2_computed = np.transpose(psi2_computed, [ 1, 2, 0 ])

    return psi2_computed
//  [python]
def testTakeAlongAxis(self, x_shape, i_shape, dtype, axis, rng_factory):
    rng = rng_factory()
    i_shape = onp.array(i_shape)
    if axis is None:
      i_shape = [onp.prod(i_shape, dtype=onp.int64)]
    else:
      # Test the case where the size of the axis doesn't necessarily broadcast.
      i_shape[axis] *= 3
      i_shape = list(i_shape)
    def args_maker():
      x = rng(x_shape, dtype)
      n = onp.prod(x_shape, dtype=onp.int32) if axis is None else x_shape[axis]
      i = rng(i_shape, onp.int32) % (2 * n - 1) - (n - 1)
      return x, i

    lnp_op = lambda x, i: tnp.take_along_axis(x, i, axis=axis)

    if hasattr(onp, "take_along_axis"):
      onp_op = lambda x, i: onp.take_along_axis(x, i, axis=axis)
      self._CheckAgainstNumpy(lnp_op, onp_op, args_maker, check_dtypes=True)
    self._CompileAndCheck(lnp_op, args_maker, check_dtypes=True,
                          check_incomplete_shape=True)
//  [python]
def sum(self, axis=None, dtype=None, out=None):
        if dtype is not None or out is not None:
            raise NotImplementedError('not implemented yet')

        if axis is None:
            tmp = n.prod(self.data.shape[2:])
            return UTPM(numpy.sum(self.data.reshape(self.data.shape[:2] + (tmp,)), axis = 2))
        else:
            if axis < 0:
                a = self.data.ndim + axis
            else:
                a = axis + 2
            return UTPM(numpy.sum(self.data, axis = a))
//  [python]
def sum(self, axis=None, dtype=None, out=None):
        if dtype is not None or out is not None:
            raise NotImplementedError('not implemented yet')

        if axis is None:
            tmp = numpy.prod(self.data.shape[2:])
            return UTPM(numpy.sum(self.data.reshape(self.data.shape[:2] + (tmp,)), axis = 2))
        else:
            if axis < 0:
                a = self.data.ndim + axis
            else:
                a = axis + 2
            return UTPM(numpy.sum(self.data, axis = a))
//  [python]
def as_utpm(cls, x):
        """ tries to convert a container (e.g. list or numpy.array) with UTPM elements as instances to a UTPM instance"""

        x_shp = numpy.shape(x)
        xr = numpy.ravel(x)

        # print 'x=', x
        # print 'xr=',xr
        # print 'x.dtype', x.dtype
        D,P = xr[0].data.shape[:2]
        shp = xr[0].data.shape[2:]

        if not isinstance(shp, tuple): shp = (shp,)
        if not isinstance(x_shp, tuple): x_shp = (x_shp,)

        y = UTPM(numpy.zeros((D,P) + x_shp + shp))

        yr = UTPM( y.data.reshape((D,P) + (reduce(lambda x,y:x*y, x_shp,1),) + shp))

        # print yr.shape
        # print yr.data.shape

        for n in range(len(xr)):
            # print yr[n].shape
            # print xr[n].shape
            yr[n] = xr[n]

        return y
//  [python]
def as_utpm(cls, x):
        """ tries to convert a container (e.g. list or numpy.array) with UTPM elements as instances to a UTPM instance"""

        x_shp = numpy.shape(x)
        xr = numpy.ravel(x)

        # print 'x=', x
        # print 'xr=',xr
        # print 'x.dtype', x.dtype
        D,P = xr[0].data.shape[:2]
        shp = xr[0].data.shape[2:]

        if not isinstance(shp, tuple): shp = (shp,)
        if not isinstance(x_shp, tuple): x_shp = (x_shp,)

        y = UTPM(numpy.zeros((D,P) + x_shp + shp))

        yr = UTPM( y.data.reshape((D,P) + (numpy.prod(x_shp, dtype=int),) + shp))

        # print yr.shape
        # print yr.data.shape

        for n in range(len(xr)):
            # print yr[n].shape
            # print xr[n].shape
            yr[n] = xr[n]

        return y
//  [python]
def _reset(self):
        """ Right now only use one initial state"""
        if hasattr(self, 'action_tool'):
            particle_pos = pyflex.get_positions().reshape(-1, 4)
            drop_point_pos = particle_pos[self._get_drop_point_idx(), :3]
            middle_point = np.mean(drop_point_pos, axis=0)
            self.action_tool.reset(middle_point)  # middle point is not really useful
            picker_radius = self.action_tool.picker_radius
            self.action_tool.update_picker_boundary([-0.3, 0.5, -0.5], [0.5, 2, 0.5])
            self.action_tool.set_picker_pos(picker_pos=drop_point_pos + np.array([0., picker_radius, 0.]))

        config = self.get_current_config()
        num_particles = np.multiply.reduce(config['ClothSize'], dtype=int)
        particle_grid_idx = np.array(list(range(num_particles))).reshape(config['ClothSize'][1], config['ClothSize'][0])  # Reversed index here

        cloth_dimx = config['ClothSize'][0]
        x_split = cloth_dimx // 2
        self.fold_group_a = particle_grid_idx[:, :x_split].flatten()
        self.fold_group_b = np.flip(particle_grid_idx, axis=1)[:, :x_split].flatten()

        colors = np.zeros(num_particles)
        colors[self.fold_group_a] = 1

        pyflex.step()
        self.init_pos = pyflex.get_positions().reshape((-1, 4))[:, :3]
        pos_a = self.init_pos[self.fold_group_a, :]
        pos_b = self.init_pos[self.fold_group_b, :]
        self.prev_dist = np.mean(np.linalg.norm(pos_a - pos_b, axis=1))

        self.performance_init = None
        info = self._get_info()
        self.performance_init = info['performance']
        return self._get_obs()
//  [python]
def _reset(self):
        """ Right now only use one initial state"""
        if hasattr(self, 'action_tool'):
            particle_pos = pyflex.get_positions().reshape(-1, 4)
            drop_point_pos = particle_pos[self._get_drop_point_idx(), :3]
            middle_point = np.mean(drop_point_pos, axis=0)
            self.action_tool.reset(middle_point)  # middle point is not really useful
            picker_radius = self.action_tool.picker_radius
            self.action_tool.update_picker_boundary([-0.3, 0.5, -0.5], [0.5, 2, 0.5])
            self.action_tool.set_picker_pos(picker_pos=drop_point_pos + np.array([0., picker_radius, 0.]))

        config = self.get_current_config()
        num_particles = np.prod(config['ClothSize'], dtype=int)
        particle_grid_idx = np.array(list(range(num_particles))).reshape(config['ClothSize'][1], config['ClothSize'][0])  # Reversed index here

        cloth_dimx = config['ClothSize'][0]
        x_split = cloth_dimx // 2
        self.fold_group_a = particle_grid_idx[:, :x_split].flatten()
        self.fold_group_b = np.flip(particle_grid_idx, axis=1)[:, :x_split].flatten()

        colors = np.zeros(num_particles)
        colors[self.fold_group_a] = 1

        pyflex.step()
        self.init_pos = pyflex.get_positions().reshape((-1, 4))[:, :3]
        pos_a = self.init_pos[self.fold_group_a, :]
        pos_b = self.init_pos[self.fold_group_b, :]
        self.prev_dist = np.mean(np.linalg.norm(pos_a - pos_b, axis=1))

        self.performance_init = None
        info = self._get_info()
        self.performance_init = info['performance']
        return self._get_obs()
//  [python]
def _numpy_prod_with_none_check(array: Union[List, None]) -> Union[float, None]:
    """
    :param array: an array like list
    :return: the product of the array if array is not None otherwise return None
    """
    return n.prod(array) if array is not None else None
//  [python]
def _numpy_prod_with_none_check(array: Union[List, None]) -> Union[float, None]:
    """
    :param array: an array like list
    :return: the product of the array if array is not None otherwise return None
    """
    return numpy.prod(array) if array is not None else None
//  [python]
def _pdf(self, **kwargs):
        """Returns the pdf at the given values. The keyword arguments must
        contain all of parameters in self's params. Unrecognized arguments are
        ignored.
        """
        for p in self._params:
            if p not in kwargs.keys():
                raise ValueError(
                            'Missing parameter {} to construct pdf.'.format(p))
        if kwargs in self:
            pdf = self._norm * \
                  reduce(lambda x, y: x*y, [(kwargs[p])**(self.dim - 1)
                              for p in self._params])
            return float(pdf)
        else:
            return 0.0
//  [python]
def _pdf(self, **kwargs):
        """Returns the pdf at the given values. The keyword arguments must
        contain all of parameters in self's params. Unrecognized arguments are
        ignored.
        """
        for p in self._params:
            if p not in kwargs.keys():
                raise ValueError(
                            'Missing parameter {} to construct pdf.'.format(p))
        if kwargs in self:
            pdf = self._norm * \
                  numpy.prod([(kwargs[p])**(self.dim - 1)
                              for p in self._params])
            return float(pdf)
        else:
            return 0.0
//  [python]
def get_num_params(self):
        """
        :return: number of model parameters, i.e. total dimension
        :rtype: int
        """
        num_params = 0
        params = self.get_params_list()
        for param in params:
            shape = param.get_shape().as_list()
            if all(shape):
                num_params += int(numpy.prod(shape))
        return num_params
//  [python]
def extract_predictions(self, response, model_components):
        score = float(n.prod(n.exp([x.logprob for x in response.generated_tokens])))
        return (Prediction(True, score) if response.generated_text.lower().strip() == self.get_pos_label(
            model_components.category_name) else Prediction(
            False, 1 - score))
//  [python]
def extract_predictions(self, response, model_components):
        score = float(numpy.prod(numpy.exp([x.logprob for x in response.generated_tokens])))
        return (Prediction(True, score) if response.generated_text.lower().strip() == self.get_pos_label(
            model_components.category_name) else Prediction(
            False, 1 - score))
//  [python]
def pb_tile(cls, Bbar, A, reps, B, out = None):

        if(isinstance(reps, int)):
            reps=[reps]

        d = len(reps)

        assert Bbar.shape == B.shape

        if A.ndim < d:
            A2shp = [1]*(d - A.ndim) + list(A.shape)
            d2shp = list(reps)
        elif A.ndim > d:
            A2shp = list(A.shape)
            d2shp = [1]*(A.ndim - d) + list(reps)
        else:
            A2shp = list(A.shape)
            d2shp = list(reps)

        if out is None:
            Abar = A.zeros_like()

        else:
            Abar = out[0]

        A2shp = numpy.array(A2shp)

        # loop over all tiles and add tile to Abar
        # K = [7*3*2, 3*2, 2, 1]
        K = [int(n.prod(d2shp[::-1][:k])) for k in range(1+len(d2shp))][::-1]
        for i in range(K[0]):
            # convert i into multi-index
            m = numpy.array([(i/K[j+1]) % d2shp[j] for j in range(len(d2shp))])
            # create slice of B
            s = [slice(A2shp[k]*m[k], A2shp[k]*(m[k]+1)) for k in range(len(m))]
            Abar += Bbar[s]

        return Abar
//  [python]
def pb_tile(cls, Bbar, A, reps, B, out = None):

        if(isinstance(reps, int)):
            reps=[reps]

        d = len(reps)

        assert Bbar.shape == B.shape

        if A.ndim < d:
            A2shp = [1]*(d - A.ndim) + list(A.shape)
            d2shp = list(reps)
        elif A.ndim > d:
            A2shp = list(A.shape)
            d2shp = [1]*(A.ndim - d) + list(reps)
        else:
            A2shp = list(A.shape)
            d2shp = list(reps)

        if out is None:
            Abar = A.zeros_like()

        else:
            Abar = out[0]

        A2shp = numpy.array(A2shp)

        # loop over all tiles and add tile to Abar
        # K = [7*3*2, 3*2, 2, 1]
        K = [int(numpy.prod(d2shp[::-1][:k])) for k in range(1+len(d2shp))][::-1]
        for i in range(K[0]):
            # convert i into multi-index
            m = numpy.array([(i/K[j+1]) % d2shp[j] for j in range(len(d2shp))])
            # create slice of B
            s = [slice(A2shp[k]*m[k], A2shp[k]*(m[k]+1)) for k in range(len(m))]
            Abar += Bbar[s]

        return Abar
//  [python]
def _calculate_flops_matmul(
    op_type: str,
    input_shape: Union[List[List], None] = None,
    output_shape: Union[List[List], None] = None,
    weight_shape: Union[List, None] = None,
) -> Union[float, None]:
    """
    Calculates flops in an ONNX MatMul operation.

    If input shape only contains 1 input, in otherwords the value of the
    first index is 1, then the matrix operation is treated as a Gemm operation.

    Otherwise the operation is treated like a NumPy operation.

    Will return none if any required value is set to None

    :param op_type: Operation type of flop calculation
    :param input_shape: List of input shapes of operation
    :param output_shape: List of output shapes of operation
    :param weight_shape: Shape of weights in operation if any, else None
    :return: The amount of floating point operations in the operation
    """
    flops = None
    if (
        input_shape is not None
        and output_shape is not None
        and len(input_shape) > 1
        and input_shape[0][-1] == input_shape[1][-2]
    ):
        matrix_ops = (
            input_shape[0][-2] * input_shape[1][-1] * (2 * input_shape[0][-1] - 1)
        )
        flops = n.prod(output_shape[0][:-2]) * matrix_ops
    elif input_shape is not None and len(input_shape) == 1:
        flops = _numpy_prod_with_none_check(weight_shape)
        flops = flops * 2 if flops is not None else None
    return flops
//  [python]
def _calculate_flops_matmul(
    op_type: str,
    input_shape: Union[List[List], None] = None,
    output_shape: Union[List[List], None] = None,
    weight_shape: Union[List, None] = None,
) -> Union[float, None]:
    """
    Calculates flops in an ONNX MatMul operation.

    If input shape only contains 1 input, in otherwords the value of the
    first index is 1, then the matrix operation is treated as a Gemm operation.

    Otherwise the operation is treated like a NumPy operation.

    Will return none if any required value is set to None

    :param op_type: Operation type of flop calculation
    :param input_shape: List of input shapes of operation
    :param output_shape: List of output shapes of operation
    :param weight_shape: Shape of weights in operation if any, else None
    :return: The amount of floating point operations in the operation
    """
    flops = None
    if (
        input_shape is not None
        and output_shape is not None
        and len(input_shape) > 1
        and input_shape[0][-1] == input_shape[1][-2]
    ):
        matrix_ops = (
            input_shape[0][-2] * input_shape[1][-1] * (2 * input_shape[0][-1] - 1)
        )
        flops = numpy.prod(output_shape[0][:-2]) * matrix_ops
    elif input_shape is not None and len(input_shape) == 1:
        flops = _numpy_prod_with_none_check(weight_shape)
        flops = flops * 2 if flops is not None else None
    return flops
//  [python]
def extract_predictions(self, response, model_component):
        score = float(n.prod(n.exp([x.logprob for x in response.generated_tokens])))
        return (Prediction(True, score) if response.generated_text.lower().strip() == self.get_pos_label(
            model_component.category_name) else Prediction(
            False, 1 - score))
//  [python]
def extract_predictions(self, response, model_component):
        score = float(numpy.prod(numpy.exp([x.logprob for x in response.generated_tokens])))
        return (Prediction(True, score) if response.generated_text.lower().strip() == self.get_pos_label(
            model_component.category_name) else Prediction(
            False, 1 - score))
//  [python]
def fanin_init(tensor):
    size = tensor.size()
    if len(size) == 2:
        fan_in = size[0]
    elif len(size) > 2:
        fan_in = n.prod(size[1:])
    else:
        raise Exception("Shape must be have dimension at least 2.")
    bound = 1. / n.sqrt(fan_in)
    return tensor.data.uniform_(-bound, bound)
//  [python]
def fanin_init(tensor):
    size = tensor.size()
    if len(size) == 2:
        fan_in = size[0]
    elif len(size) > 2:
        fan_in = np.prod(size[1:])
    else:
        raise Exception("Shape must be have dimension at least 2.")
    bound = 1. / np.sqrt(fan_in)
    return tensor.data.uniform_(-bound, bound)
//  [python]
def build_grid(space: Space, n_values: dict[str, int], max_trials: int = 10000):
        """Build a grid of trials

        Parameters
        ----------
        n_values: int or dict
            Dictionary specifying number of trials for each dimension independently
            (name, n_values). For categorical dimensions, n_values will not be used, and all
            categories will be used to build the grid.
        max_trials: int
            Maximum number of trials for the grid. If n_values lead to more trials than max_trials,
            the n_values will be adjusted down. Will raise ValueError if it is impossible to build
            a grid smaller than max_trials (for instance if choices are too large).

        """
        adjust = 0
        n_trials = float("inf")
        coordinates: list[list] = []

        while n_trials > max_trials:
            coordinates = []
            capped_values = []
            for name, dim in space.items():
                capped_value = max(n_values[name] - adjust, 1)
                capped_values.append(capped_value)
                coordinates.append(list(grid(dim, capped_value)))

            if all(value <= 1 for value in capped_values):
                raise ValueError(
                    f"Cannot build a grid smaller than {max_trials}. "
                    "Try reducing the number of choices in categorical dimensions."
                )

            n_trials = n.prod([len(dim_values) for dim_values in coordinates])
            # TODO: Use binary search instead of incrementing by one.
            adjust += 1

        if adjust > 1:
            log.warning(
                f"`n_values` reduced by {adjust-1} to limit number of trials below {max_trials}."
            )

        return list(itertools.product(*coordinates))
//  [python]
def build_grid(space: Space, n_values: dict[str, int], max_trials: int = 10000):
        """Build a grid of trials

        Parameters
        ----------
        n_values: int or dict
            Dictionary specifying number of trials for each dimension independently
            (name, n_values). For categorical dimensions, n_values will not be used, and all
            categories will be used to build the grid.
        max_trials: int
            Maximum number of trials for the grid. If n_values lead to more trials than max_trials,
            the n_values will be adjusted down. Will raise ValueError if it is impossible to build
            a grid smaller than max_trials (for instance if choices are too large).

        """
        adjust = 0
        n_trials = float("inf")
        coordinates: list[list] = []

        while n_trials > max_trials:
            coordinates = []
            capped_values = []
            for name, dim in space.items():
                capped_value = max(n_values[name] - adjust, 1)
                capped_values.append(capped_value)
                coordinates.append(list(grid(dim, capped_value)))

            if all(value <= 1 for value in capped_values):
                raise ValueError(
                    f"Cannot build a grid smaller than {max_trials}. "
                    "Try reducing the number of choices in categorical dimensions."
                )

            n_trials = numpy.prod([len(dim_values) for dim_values in coordinates])
            # TODO: Use binary search instead of incrementing by one.
            adjust += 1

        if adjust > 1:
            log.warning(
                f"`n_values` reduced by {adjust-1} to limit number of trials below {max_trials}."
            )

        return list(itertools.product(*coordinates))
//  [python]
def _is_f_contiguous(shape, strides, itemsize):
    if n.prod(shape) <= 1:
        return True
    for sh, st in zip(shape, reversed(strides)):
        if sh == 1:
            continue
        if st != itemsize:
            return False
        itemsize *= sh
    return True
//  [python]
def _is_f_contiguous(shape, strides, itemsize):
    if numpy.prod(shape) <= 1:
        return True
    for sh, st in zip(shape, reversed(strides)):
        if sh == 1:
            continue
        if st != itemsize:
            return False
        itemsize *= sh
    return True
//  [python]
def fanin_init_weights_like(tensor):
    size = tensor.size()
    if len(size) == 2:
        fan_in = size[0]
    elif len(size) > 2:
        fan_in = n.prod(size[1:])
    else:
        raise Exception("Shape must be have dimension at least 2.")
    bound = 1. / n.sqrt(fan_in)
    new_tensor = FloatTensor(tensor.size())
    new_tensor.uniform_(-bound, bound)
    return new_tensor
//  [python]
def fanin_init_weights_like(tensor):
    size = tensor.size()
    if len(size) == 2:
        fan_in = size[0]
    elif len(size) > 2:
        fan_in = np.prod(size[1:])
    else:
        raise Exception("Shape must be have dimension at least 2.")
    bound = 1. / np.sqrt(fan_in)
    new_tensor = FloatTensor(tensor.size())
    new_tensor.uniform_(-bound, bound)
    return new_tensor
//  [python]
def extract_predictions(self, response,  model_components):
        category_names = list(model_components.category_name_to_id.keys())
        score = float(n.prod(n.exp([x.logprob for x in response.generated_tokens])))
        # in case prediction is not in the classes list:
        predicted_class = response.generated_text.lower().strip()
        if predicted_class not in category_names:
            closest_class = difflib.get_close_matches(response.generated_text, category_names, n=1)
            if (len(closest_class)) > 0:
                predicted_class = closest_class[0]
            else:
                if model_components.sorted_categories_by_freq and len(model_components.sorted_categories_by_freq) > 0:
                    #this is initialized during tuning
                        predicted_class = model_components.sorted_categories_by_freq[0][0]

        predicted_label = model_components.category_name_to_id.get(predicted_class)
        if not predicted_label:
            #default label is the first in the categories list
            predicted_label = model_components.category_name_to_id.get(category_names[0])

        scores = {i: 0 for i in model_components.category_name_to_id.values()}
        scores[predicted_label] = score
        return MulticlassPrediction(label=predicted_label, scores=scores)
//  [python]
def extract_predictions(self, response,  model_components):
        category_names = list(model_components.category_name_to_id.keys())
        score = float(numpy.prod(numpy.exp([x.logprob for x in response.generated_tokens])))
        # in case prediction is not in the classes list:
        predicted_class = response.generated_text.lower().strip()
        if predicted_class not in category_names:
            closest_class = difflib.get_close_matches(response.generated_text, category_names, n=1)
            if (len(closest_class)) > 0:
                predicted_class = closest_class[0]
            else:
                if model_components.sorted_categories_by_freq and len(model_components.sorted_categories_by_freq) > 0:
                    #this is initialized during tuning
                        predicted_class = model_components.sorted_categories_by_freq[0][0]

        predicted_label = model_components.category_name_to_id.get(predicted_class)
        if not predicted_label:
            #default label is the first in the categories list
            predicted_label = model_components.category_name_to_id.get(category_names[0])

        scores = {i: 0 for i in model_components.category_name_to_id.values()}
        scores[predicted_label] = score
        return MulticlassPrediction(label=predicted_label, scores=scores)