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

problem_size = (4096, 4096)
size = numpy.prod(problem_size)
//  [python]
def _push_allocation_config(self):
        self.conv_sequence._push_allocation_config()
        conv_out_dim = self.conv_sequence.get_dim('output')

        self.top_mlp.activations = self.top_mlp_activations
        self.top_mlp.dims = [n.prod(conv_out_dim)] + self.top_mlp_dims
//  [python]
def _push_allocation_config(self):
        self.conv_sequence._push_allocation_config()
        conv_out_dim = self.conv_sequence.get_dim('output')

        self.top_mlp.activations = self.top_mlp_activations
        self.top_mlp.dims = [numpy.prod(conv_out_dim)] + self.top_mlp_dims
//  [python]
def size(self):
        r"""Returns the size of the self :class:`~.Tensor`.
        The returned value is a subclass of :class:`tuple`.
        """
        shape = self.shape
        if shape.__class__ is tuple:
            return np.prod(self.shape).item()
        return shape.prod()
//  [python]
def _loadtxt_flatten_dtype_internal(dt):
    """Unpack a structured data-type, and produce a packer function."""
    if dt.names is None:
        # If the dtype is flattened, return.
        # If the dtype has a shape, the dtype occurs
        # in the list more than once.
        shape = dt.shape
        if len(shape) == 0:
            return ([dt.base], None)
        else:
            packing = [(shape[-1], list)]
            if len(shape) > 1:
                for dim in dt.shape[-2::-1]:
                    packing = [(dim*packing[0][0], packing*dim)]
            return ([dt.base] * int(n.prod(dt.shape)),
                    functools.partial(_loadtxt_pack_items, packing))
    else:
        types = []
        packing = []
        for field in dt.names:
            tp, bytes = dt.fields[field]
            flat_dt, flat_packer = _loadtxt_flatten_dtype_internal(tp)
            types.extend(flat_dt)
            flat_packing = flat_packer.args[0] if flat_packer else None
            # Avoid extra nesting for subarrays
            if tp.ndim > 0:
                packing.extend(flat_packing)
            else:
                packing.append((len(flat_dt), flat_packing))
        return (types, functools.partial(_loadtxt_pack_items, packing))
//  [python]
def _loadtxt_flatten_dtype_internal(dt):
    """Unpack a structured data-type, and produce a packer function."""
    if dt.names is None:
        # If the dtype is flattened, return.
        # If the dtype has a shape, the dtype occurs
        # in the list more than once.
        shape = dt.shape
        if len(shape) == 0:
            return ([dt.base], None)
        else:
            packing = [(shape[-1], list)]
            if len(shape) > 1:
                for dim in dt.shape[-2::-1]:
                    packing = [(dim*packing[0][0], packing*dim)]
            return ([dt.base] * int(np.prod(dt.shape)),
                    functools.partial(_loadtxt_pack_items, packing))
    else:
        types = []
        packing = []
        for field in dt.names:
            tp, bytes = dt.fields[field]
            flat_dt, flat_packer = _loadtxt_flatten_dtype_internal(tp)
            types.extend(flat_dt)
            flat_packing = flat_packer.args[0] if flat_packer else None
            # Avoid extra nesting for subarrays
            if tp.ndim > 0:
                packing.extend(flat_packing)
            else:
                packing.append((len(flat_dt), flat_packing))
        return (types, functools.partial(_loadtxt_pack_items, packing))
//  [python]
def tune():
    with open('convolution.cu', 'r') as f:
        kernel_string = f.read()

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

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

    #limit the search to only use padding when its effective, and at least 32 threads in a block
    restrict = ["use_padding==0 or (block_size_x % 32 != 0)", "block_size_x*block_size_y >= 32"]

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

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

    #setup kernel arguments
    cmem_args = {'d_filter': filter_weights}
    args = [output_image, input_image, filter_weights]

    #tell the Kernel Tuner how to compute grid dimensions
    grid_div_x = ["block_size_x", "tile_size_x"]
    grid_div_y = ["block_size_y", "tile_size_y"]

    #start tuning separable convolution (row)
    tune_params["filter_height"] = [1]
    tune_params["tile_size_y"] = [1]
    results_row = tune_kernel("convolution_kernel", kernel_string,
        problem_size, args, tune_params,
        grid_div_y=grid_div_y, grid_div_x=grid_div_x, cmem_args=cmem_args, verbose=True, restrictions=restrict)

    #start tuning separable convolution (col)
    tune_params["filter_height"] = tune_params["filter_width"][:]
    tune_params["file_size_y"] = tune_params["tile_size_x"][:]
    tune_params["filter_width"] = [1]
    tune_params["tile_size_x"] = [1]
    results_col = tune_kernel("convolution_kernel", kernel_string,
        problem_size, args, tune_params,
        grid_div_y=grid_div_y, grid_div_x=grid_div_x, cmem_args=cmem_args, verbose=True, restrictions=restrict)

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

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

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

    #limit the search to only use padding when its effective, and at least 32 threads in a block
    restrict = ["use_padding==0 or (block_size_x % 32 != 0)", "block_size_x*block_size_y >= 32"]

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

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

    #setup kernel arguments
    cmem_args = {'d_filter': filter_weights}
    args = [output_image, input_image, filter_weights]

    #tell the Kernel Tuner how to compute grid dimensions
    grid_div_x = ["block_size_x", "tile_size_x"]
    grid_div_y = ["block_size_y", "tile_size_y"]

    #start tuning separable convolution (row)
    tune_params["filter_height"] = [1]
    tune_params["tile_size_y"] = [1]
    results_row = tune_kernel("convolution_kernel", kernel_string,
        problem_size, args, tune_params,
        grid_div_y=grid_div_y, grid_div_x=grid_div_x, cmem_args=cmem_args, verbose=True, restrictions=restrict)

    #start tuning separable convolution (col)
    tune_params["filter_height"] = tune_params["filter_width"][:]
    tune_params["file_size_y"] = tune_params["tile_size_x"][:]
    tune_params["filter_width"] = [1]
    tune_params["tile_size_x"] = [1]
    results_col = tune_kernel("convolution_kernel", kernel_string,
        problem_size, args, tune_params,
        grid_div_y=grid_div_y, grid_div_x=grid_div_x, cmem_args=cmem_args, verbose=True, restrictions=restrict)

    return results_row, results_col
//  [python]
def cartesian(arrays, out=None):
        """
        Generate a cartesian product of input arrays.

        Parameters
        ----------
        arrays : list of array-like
                1-D arrays to form the cartesian product of.
        out : ndarray
                Array to place the cartesian product in.

        Returns
        -------
        out : ndarray
                2-D array of shape (M, len(arrays)) containing cartesian products
                formed of input arrays.

        Examples
        --------
        >>> cartesian(([1, 2, 3], [4, 5], [6, 7]))
        array([[1, 4, 6],
                     [1, 4, 7],
                     [1, 5, 6],
                     [1, 5, 7],
                     [2, 4, 6],
                     [2, 4, 7],
                     [2, 5, 6],
                     [2, 5, 7],
                     [3, 4, 6],
                     [3, 4, 7],
                     [3, 5, 6],
                     [3, 5, 7]])

        """

        arrays = [np.asarray(x) for x in arrays]
        dtype = arrays[0].dtype

        n = 1
        for x in arrays:
            n *= x.size
        if out is None:
            out = np.zeros([n, len(arrays)], dtype=dtype)

        m = n / arrays[0].size
        m = int(m)
        out[:, 0] = np.repeat(arrays[0], m)
        if arrays[1:]:
            cartesian(arrays[1:], out=out[0:m, 1:])
            for j in range(1, arrays[0].size):
                out[j * m:(j + 1) * m, 1:] = out[0:m, 1:]
        return out
//  [python]
def cartesian(arrays, out=None):
        """
        Generate a cartesian product of input arrays.

        Parameters
        ----------
        arrays : list of array-like
                1-D arrays to form the cartesian product of.
        out : ndarray
                Array to place the cartesian product in.

        Returns
        -------
        out : ndarray
                2-D array of shape (M, len(arrays)) containing cartesian products
                formed of input arrays.

        Examples
        --------
        >>> cartesian(([1, 2, 3], [4, 5], [6, 7]))
        array([[1, 4, 6],
                     [1, 4, 7],
                     [1, 5, 6],
                     [1, 5, 7],
                     [2, 4, 6],
                     [2, 4, 7],
                     [2, 5, 6],
                     [2, 5, 7],
                     [3, 4, 6],
                     [3, 4, 7],
                     [3, 5, 6],
                     [3, 5, 7]])

        """

        arrays = [np.asarray(x) for x in arrays]
        dtype = arrays[0].dtype

        n = np.prod([x.size for x in arrays])
        if out is None:
            out = np.zeros([n, len(arrays)], dtype=dtype)

        m = n / arrays[0].size
        m = int(m)
        out[:, 0] = np.repeat(arrays[0], m)
        if arrays[1:]:
            cartesian(arrays[1:], out=out[0:m, 1:])
            for j in range(1, arrays[0].size):
                out[j * m:(j + 1) * m, 1:] = out[0:m, 1:]
        return out
//  [python]
def count_network_parameters(model):
    parameters = filter(lambda p: p.requires_grad, model.parameters())
    N = sum([n.prod(p.size()) for p in parameters])
    return N
//  [python]
def count_network_parameters(model):
    parameters = filter(lambda p: p.requires_grad, model.parameters())
    N = sum([numpy.prod(p.size()) for p in parameters])
    return N
//  [python]
def to_hdf5(self, h5group):
        # Write Datasets
        for key in self.keys():
            if self.__getattribute__(key) is None:
                continue
            if hasattr(self.__getattribute__(key), 'to_hdf5'):
                subgroup = h5group.create_group(key)
                self.__getattribute__(key).to_hdf5(subgroup)
            elif self.__getattribute__(key).dtype == 'object':
                # Handle numpy type = object as vlen string
                h5group.create_dataset(key,
                                       data=self.__getattribute__(
                                           key).tolist().__repr__(),
                                       dtype=h5py.special_dtype(vlen=str))
            else:
                if np.prod(self.__getattribute__(key).shape):
                    maxshape = None
                else:
                    maxshape = (None,)
                h5group.create_dataset(
                    key, data=self.__getattribute__(key), maxshape=maxshape)
//  [python]
def test_flatten(self, array):
        flat = array.flatten()
        assert flat.ndim == 1
        assert type(flat) is units.Quantity  # pylint: disable=C0123
        assert flat.shape[0] == n.prod(array.shape)
        try:
            utils.assert_quantity_equal(
                array.flatten('C'),
                array.T.flatten('F'),
            )
        except NotImplementedError:
            pass
//  [python]
def test_flatten(self, array):
        flat = array.flatten()
        assert flat.ndim == 1
        assert type(flat) is units.Quantity  # pylint: disable=C0123
        assert flat.shape[0] == numpy.prod(array.shape)
        try:
            utils.assert_quantity_equal(
                array.flatten('C'),
                array.T.flatten('F'),
            )
        except NotImplementedError:
            pass
//  [python]
def __call__(self, x):
        if self.mu.W.array is None:
            self.mu.W.initialize((self.out_size, n.prod(x.shape[1:])))
        if self.sigma.W.array is None:
            self.sigma.W.initialize((self.out_size, n.prod(x.shape[1:])))

        # use info of sigma.W to avoid strange error messages
        dtype = self.sigma.W.dtype
        out_size, in_size = self.sigma.W.shape

        eps = self._eps(in_size + out_size, dtype)
        eps_x = eps[:in_size]
        eps_y = eps[in_size:]
        W = muladd(self.sigma.W, self.xp.outer(eps_y, eps_x), self.mu.W)
        if self.nobias:
            return F.linear(x, W)
        else:
            b = muladd(self.sigma.b, eps_y, self.mu.b)
            return F.linear(x, W, b)
//  [python]
def __call__(self, x):
        if self.mu.W.array is None:
            self.mu.W.initialize((self.out_size, numpy.prod(x.shape[1:])))
        if self.sigma.W.array is None:
            self.sigma.W.initialize((self.out_size, numpy.prod(x.shape[1:])))

        # use info of sigma.W to avoid strange error messages
        dtype = self.sigma.W.dtype
        out_size, in_size = self.sigma.W.shape

        eps = self._eps(in_size + out_size, dtype)
        eps_x = eps[:in_size]
        eps_y = eps[in_size:]
        W = muladd(self.sigma.W, self.xp.outer(eps_y, eps_x), self.mu.W)
        if self.nobias:
            return F.linear(x, W)
        else:
            b = muladd(self.sigma.b, eps_y, self.mu.b)
            return F.linear(x, W, b)
//  [python]
def agent_init(self,taskSpec):
        sarsa_lambda.sarsa_lambda.agent_init(self, taskSpec)
        self.traces = numpy.zeros((numpy.prod(self.weights.shape[:-1]) + self.weights.size,))
        self.A = numpy.eye(self.traces.size)*self.nac_precond
#        self.A += numpy.random.random(self.A.shape)*self.nac_precond
        self.b = numpy.zeros((self.traces.size,))
//  [python]
def testPropagatesProcessingException(self, op_raising_at_3):
        roiList = [
            ((0, 0, 0), (4, 4, 4)),
            ((1, 1, 1), (4, 4, 4)),
            ((2, 2, 2), (4, 4, 4)),
            ((3, 3, 3), (4, 4, 4)),
        ]

        totalVolume = n.prod(op_raising_at_3.Output.meta.shape)
        batch = RoiRequestBatch(
            op_raising_at_3.Output, roiList.__iter__(), totalVolume, batchSize=1, allowParallelResults=False
        )

        with pytest.raises(RoiRequestBatchException) as exc_info:
            batch.execute()

        assert is_root_cause(ProcessingException, exc_info.value)
//  [python]
def testPropagatesProcessingException(self, op_raising_at_3):
        roiList = [
            ((0, 0, 0), (4, 4, 4)),
            ((1, 1, 1), (4, 4, 4)),
            ((2, 2, 2), (4, 4, 4)),
            ((3, 3, 3), (4, 4, 4)),
        ]

        totalVolume = numpy.prod(op_raising_at_3.Output.meta.shape)
        batch = RoiRequestBatch(
            op_raising_at_3.Output, roiList.__iter__(), totalVolume, batchSize=1, allowParallelResults=False
        )

        with pytest.raises(RoiRequestBatchException) as exc_info:
            batch.execute()

        assert is_root_cause(ProcessingException, exc_info.value)
//  [python]
def testFailedResultHandler(self):
        op = OpArrayPiper(graph=Graph())
        inputData = numpy.indices((100, 100)).sum(0)
        op.Input.setValue(inputData)
        roiList = []
        block_starts = getIntersectingBlocks([10, 10], ([0, 0], [100, 100]))
        for block_start in block_starts:
            roiList.append(getBlockBounds([100, 100], [10, 10], block_start))

        class SpecialException(Exception):
            pass

        def handleResult(roi, result):
            raise SpecialException("Intentional Exception: raised while handling the result")

        totalVolume = numpy.prod(inputData.shape)
        batch = RoiRequestBatch(op.Output, roiList.__iter__(), totalVolume, batchSize=10, allowParallelResults=False)
        batch.resultSignal.subscribe(handleResult)

        # FIXME: There are multiple places where the RoiRequestBatch tool should be prepared to handle exceptions.
        #        This only tests one of them (in the notify_finished() handler)
        with pytest.raises(RoiRequestBatchException) as exc_info:
            batch.execute()

        assert is_root_cause(SpecialException, exc_info.value)
//  [python]
def agent_init(self,taskSpec):
        sarsa_lambda.sarsa_lambda.agent_init(self, taskSpec)
        self.traces = numpy.zeros((n.prod(self.weights.shape[:-1]) + self.weights.size,)) # combined e_t^w and e_t^v
        self.value_weights = numpy.zeros((n.prod(self.weights.shape[:-1]),))
        self.advantage_weights = numpy.zeros((self.weights.size,))
//  [python]
def agent_init(self,taskSpec):
        sarsa_lambda.sarsa_lambda.agent_init(self, taskSpec)
        self.traces = numpy.zeros((numpy.prod(self.weights.shape[:-1]) + self.weights.size,)) # combined e_t^w and e_t^v
        self.value_weights = numpy.zeros((numpy.prod(self.weights.shape[:-1]),))
        self.advantage_weights = numpy.zeros((self.weights.size,))
//  [python]
def testBasic(self):
        op = OpArrayPiper(graph=Graph())
        inputData = numpy.indices((100, 100)).sum(0)
        op.Input.setValue(inputData)
        roiList = []
        block_starts = getIntersectingBlocks([10, 10], ([0, 0], [100, 100]))
        for block_start in block_starts:
            roiList.append(getBlockBounds([100, 100], [10, 10], block_start))

        results = numpy.zeros((100, 100), dtype=numpy.int32)
        resultslock = threading.Lock()

        resultsCount = [0]

        def handleResult(roi, result):
            acquired = resultslock.acquire(False)
            assert acquired, "resultslock is contested! Access to callback is supposed to be automatically serialized."
            results[roiToSlice(*roi)] = result
            logger.debug("Got result for {}".format(roi))
            resultslock.release()
            resultsCount[0] += 1

        progressList = []

        def handleProgress(progress):
            progressList.append(progress)
            logger.debug("Progress update: {}".format(progress))

        totalVolume = numpy.prod(inputData.shape)
        batch = RoiRequestBatch(op.Output, roiList.__iter__(), totalVolume, batchSize=10, allowParallelResults=False)
        batch.resultSignal.subscribe(handleResult)
        batch.progressSignal.subscribe(handleProgress)

        batch.execute()
        logger.debug("Got {} results".format(resultsCount[0]))
        assert (results == inputData).all()

        # Progress reporting MUST start with 0 and end with 100
        assert progressList[0] == 0, "Invalid progress reporting."
        assert progressList[-1] == 100, "Invalid progress reporting."

        # There should be some intermediate progress reporting, but exactly how much is unspecified.
        assert len(progressList) >= 10

        logger.debug("FINISHED")
//  [python]
def candidate_weighting(self, window=2, stoplist=None, use_stems=False):
        """Candidate weight calculation as described in the YAKE paper.

        Args:
            stoplist (list): the stoplist for filtering candidates, defaults to
                the nltk stoplist.
            use_stems (bool): whether to use stems instead of lowercase words
                for weighting, defaults to False.
            window (int): the size in words of the window used for computing
                co-occurrence counts, defaults to 2.
        """
        # change: Consider stopwords between non-stopwords on condidate weight generation

        # build the vocabulary
        self._vocabulary_building(use_stems=use_stems)

        # extract the contexts
        self._contexts_building(use_stems=use_stems, window=window)

        # compute the word features
        self._feature_extraction(stoplist=stoplist)

        # compute candidate weights
        for k, v in self.candidates.items():

            # use stems
            if use_stems:
                weights = [self.features[t]['weight'] for t in v.lexical_form]
                self.weights[k] = n.prod(weights)
                self.weights[k] /= len(v.offsets) * (1 + sum(weights))

            # use words
            else:
                lowercase_forms = [' '.join(t).lower() for t in v.surface_forms]
                for i, candidate in enumerate(lowercase_forms):
                    TF = lowercase_forms.count(candidate)
                    words_cand = [t.lower() for t in v.surface_forms[i]]
                    prod_ = 1.
                    sum_  = 0.
                    for j, word_cand in enumerate(words_cand):
                        if self.features[word_cand]['isstop']:
                            """
                                self.contexts[word][0].append(w) # left
                                self.contexts[w][1].append(word) # right
                            """
                            term_left  = words_cand[j-1]
                            term_right = words_cand[j+1]
                            term_stop  = word_cand

                            prob_t1 = self.contexts[term_left][1].count(term_stop) / self.features[term_left]['TF']

                            prob_t2 = self.contexts[term_stop][0].count(term_right) / self.features[term_right]['TF']

                            prob = prob_t1 * prob_t2
                            prod_ *= (1 + (1 - prob ) )
                            sum_ -= (1 - prob)
                        else:
                            prod_ *= self.features[word_cand]['weight']
                            sum_  += self.features[word_cand]['weight']
                    self.weights[candidate] = prod_
                    self.weights[candidate] /= TF * (1 + sum_)
                    self.surface_to_lexical[candidate] = k
//  [python]
def candidate_weighting(self, window=2, stoplist=None, use_stems=False):
        """Candidate weight calculation as described in the YAKE paper.

        Args:
            stoplist (list): the stoplist for filtering candidates, defaults to
                the nltk stoplist.
            use_stems (bool): whether to use stems instead of lowercase words
                for weighting, defaults to False.
            window (int): the size in words of the window used for computing
                co-occurrence counts, defaults to 2.
        """
        # change: Consider stopwords between non-stopwords on condidate weight generation

        # build the vocabulary
        self._vocabulary_building(use_stems=use_stems)

        # extract the contexts
        self._contexts_building(use_stems=use_stems, window=window)

        # compute the word features
        self._feature_extraction(stoplist=stoplist)

        # compute candidate weights
        for k, v in self.candidates.items():

            # use stems
            if use_stems:
                weights = [self.features[t]['weight'] for t in v.lexical_form]
                self.weights[k] = numpy.prod(weights)
                self.weights[k] /= len(v.offsets) * (1 + sum(weights))

            # use words
            else:
                lowercase_forms = [' '.join(t).lower() for t in v.surface_forms]
                for i, candidate in enumerate(lowercase_forms):
                    TF = lowercase_forms.count(candidate)
                    words_cand = [t.lower() for t in v.surface_forms[i]]
                    prod_ = 1.
                    sum_  = 0.
                    for j, word_cand in enumerate(words_cand):
                        if self.features[word_cand]['isstop']:
                            """
                                self.contexts[word][0].append(w) # left
                                self.contexts[w][1].append(word) # right
                            """
                            term_left  = words_cand[j-1]
                            term_right = words_cand[j+1]
                            term_stop  = word_cand

                            prob_t1 = self.contexts[term_left][1].count(term_stop) / self.features[term_left]['TF']

                            prob_t2 = self.contexts[term_stop][0].count(term_right) / self.features[term_right]['TF']

                            prob = prob_t1 * prob_t2
                            prod_ *= (1 + (1 - prob ) )
                            sum_ -= (1 - prob)
                        else:
                            prod_ *= self.features[word_cand]['weight']
                            sum_  += self.features[word_cand]['weight']
                    self.weights[candidate] = prod_
                    self.weights[candidate] /= TF * (1 + sum_)
                    self.surface_to_lexical[candidate] = k
//  [python]
def reference_native_layer_norm(inp: np.ndarray, normalized_shape: Tuple[int], weight, bias, eps):
    feature_size = np.multiply.reduce(normalized_shape)
    inp_view = inp.reshape(-1, feature_size)  # type: ignore[call-overload]
    mean = inp_view.mean(axis=-1, keepdims=True)
    var = inp_view.var(axis=-1, ddof=0, keepdims=True)
    Y = (inp_view - mean) / np.sqrt(var + eps)
    if weight is None and bias is not None:
        Y = Y + bias.reshape(-1)
    elif weight is not None and bias is None:
        Y = Y * weight.reshape(-1)
    elif weight is not None and bias is not None:
        Y = Y * weight.reshape(-1) + bias.reshape(-1)
    axis = inp.ndim - len(normalized_shape)
    stat_shape = inp.shape[:axis] + (1,) * len(normalized_shape)
    return Y.reshape(*inp.shape), mean.reshape(stat_shape), (1.0 / np.sqrt(var + eps)).reshape(stat_shape)
//  [python]
def reference_native_layer_norm(inp: np.ndarray, normalized_shape: Tuple[int], weight, bias, eps):
    feature_size = np.prod(normalized_shape)
    inp_view = inp.reshape(-1, feature_size)  # type: ignore[call-overload]
    mean = inp_view.mean(axis=-1, keepdims=True)
    var = inp_view.var(axis=-1, ddof=0, keepdims=True)
    Y = (inp_view - mean) / np.sqrt(var + eps)
    if weight is None and bias is not None:
        Y = Y + bias.reshape(-1)
    elif weight is not None and bias is None:
        Y = Y * weight.reshape(-1)
    elif weight is not None and bias is not None:
        Y = Y * weight.reshape(-1) + bias.reshape(-1)
    axis = inp.ndim - len(normalized_shape)
    stat_shape = inp.shape[:axis] + (1,) * len(normalized_shape)
    return Y.reshape(*inp.shape), mean.reshape(stat_shape), (1.0 / np.sqrt(var + eps)).reshape(stat_shape)
//  [python]
def reference_group_norm(inp: np.ndarray, num_groups: int, weight=None, bias=None, eps=1e-5):
    inp_view = inp
    if n.prod(inp.shape) != 0:
        inp_view = inp.reshape((inp.shape[0], num_groups, -1))
    mean = inp_view.mean(axis=-1, keepdims=True)
    var = inp_view.var(axis=-1, ddof=0, keepdims=True)
    Y = (inp_view - mean) / np.sqrt(var + eps)
    Y = Y.reshape(inp.shape)
    if weight is not None:
        # weight is a vector of length equal to the channel
        if len(Y.shape) > 2:
            weight = np.expand_dims(weight, [0] + [idx + 2 for idx in range(inp.ndim - 2)])
        Y = Y * weight
    if bias is not None:
        # bias is a vector of length equal to the channel
        if len(Y.shape) > 2:
            bias = np.expand_dims(bias, [0] + [idx + 2 for idx in range(inp.ndim - 2)])
        Y = Y + bias
    return Y
//  [python]
def reference_group_norm(inp: np.ndarray, num_groups: int, weight=None, bias=None, eps=1e-5):
    inp_view = inp
    if np.prod(inp.shape) != 0:
        inp_view = inp.reshape((inp.shape[0], num_groups, -1))
    mean = inp_view.mean(axis=-1, keepdims=True)
    var = inp_view.var(axis=-1, ddof=0, keepdims=True)
    Y = (inp_view - mean) / np.sqrt(var + eps)
    Y = Y.reshape(inp.shape)
    if weight is not None:
        # weight is a vector of length equal to the channel
        if len(Y.shape) > 2:
            weight = np.expand_dims(weight, [0] + [idx + 2 for idx in range(inp.ndim - 2)])
        Y = Y * weight
    if bias is not None:
        # bias is a vector of length equal to the channel
        if len(Y.shape) > 2:
            bias = np.expand_dims(bias, [0] + [idx + 2 for idx in range(inp.ndim - 2)])
        Y = Y + bias
    return Y
//  [python]
def makeTaskSpec(self):
        ts = TaskSpecRLGlue.TaskSpec(discount_factor=1.0, reward_range=(-1.0, 0.0))
        ts.setDiscountFactor(1.0)
        ts.addDiscreteAction((0, self.state.shape[0]*2))
        flattened_ranges = self.state_range.reshape((numpy.prod(self.state_range.shape[:2]), self.state_range.shape[2]))
        for minValue, maxValue in flattened_ranges:
            ts.addContinuousObservation((minValue, maxValue))
        ts.setEpisodic()
        ts.setExtra(self.name)
        return ts.toTaskSpec()
//  [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 size(self):
        """Numpy-style attribute giving the total dataset size"""
        if self._shape is None:
            return None
        return numpy.prod(self._shape, dtype=numpy.int64).item()
//  [python]
def reference_searchsorted(sorted_sequence, boundary, out_int32=False, right=False, side='left', sorter=None):
    side = 'right' if (right or side == 'right') else 'left'
    if len(sorted_sequence.shape) == 1 :
        ret = np.searchsorted(sorted_sequence, boundary, side=side, sorter=sorter)
        return ret.astype(np.int32) if out_int32 else ret
    elif sorted_sequence.shape[0] == 0:
        if sorter is not None:
            sorter = sorter.flatten()
        ret = np.searchsorted(sorted_sequence.flatten(), boundary.flatten(), side=side, sorter=sorter)
        ret = ret.astype(np.int32) if out_int32 else ret
        return ret.reshape(boundary.shape)
    else:
        # numpy searchsorted only supports 1D inputs so we split up ND inputs
        orig_shape = boundary.shape
        num_splits = np.multiply.reduce(sorted_sequence.shape[:-1])
        splits = range(0, num_splits)
        sorted_sequence, boundary = sorted_sequence.reshape(num_splits, -1), boundary.reshape(num_splits, -1)
        if sorter is not None:
            sorter = sorter.reshape(num_splits, -1)

        split_sequence = [sorted_sequence[i] for i in splits]
        split_boundary = [boundary[i] for i in splits]
        split_sorter = [sorter[i] if (sorter is not None) else None for i in splits]

        split_ret = [np.searchsorted(s_seq, b, side=side, sorter=s_sort)
                     for (s_seq, b, s_sort) in zip(split_sequence, split_boundary, split_sorter)]
        split_ret = [i.astype(np.int32) for i in split_ret] if out_int32 else split_ret
        return np.stack(split_ret).reshape(orig_shape)
//  [python]
def reference_searchsorted(sorted_sequence, boundary, out_int32=False, right=False, side='left', sorter=None):
    side = 'right' if (right or side == 'right') else 'left'
    if len(sorted_sequence.shape) == 1 :
        ret = np.searchsorted(sorted_sequence, boundary, side=side, sorter=sorter)
        return ret.astype(np.int32) if out_int32 else ret
    elif sorted_sequence.shape[0] == 0:
        if sorter is not None:
            sorter = sorter.flatten()
        ret = np.searchsorted(sorted_sequence.flatten(), boundary.flatten(), side=side, sorter=sorter)
        ret = ret.astype(np.int32) if out_int32 else ret
        return ret.reshape(boundary.shape)
    else:
        # numpy searchsorted only supports 1D inputs so we split up ND inputs
        orig_shape = boundary.shape
        num_splits = np.prod(sorted_sequence.shape[:-1])
        splits = range(0, num_splits)
        sorted_sequence, boundary = sorted_sequence.reshape(num_splits, -1), boundary.reshape(num_splits, -1)
        if sorter is not None:
            sorter = sorter.reshape(num_splits, -1)

        split_sequence = [sorted_sequence[i] for i in splits]
        split_boundary = [boundary[i] for i in splits]
        split_sorter = [sorter[i] if (sorter is not None) else None for i in splits]

        split_ret = [np.searchsorted(s_seq, b, side=side, sorter=s_sort)
                     for (s_seq, b, s_sort) in zip(split_sequence, split_boundary, split_sorter)]
        split_ret = [i.astype(np.int32) for i in split_ret] if out_int32 else split_ret
        return np.stack(split_ret).reshape(orig_shape)
//  [python]
def __array__(self, dtype=None):
        """Create a Numpy array containing the whole dataset.  DON'T THINK
        THIS MEANS DATASETS ARE INTERCHANGABLE WITH ARRAYS.  For one thing,
        you have to read the whole dataset everytime this method is called.
        """
        arr = numpy.empty(self._shape, dtype=self.dtype if dtype is None else dtype)

        # Special case for (0,)*-shape datasets
        if self._shape is None or numpy.prod(self._shape) == 0:
            return arr

        self.read_direct(arr)
        return arr
//  [python]
a = np.array([[1,2,3],[4,5,6]])
print(numpy.size(a))
print(numpy.size(a,1))
print(numpy.size(a,0))

print(np.prod([1, 10, 100, 5]))
//  [python]
a = np.array([[1,2,3],[4,5,6]])
print(numpy.size(a))
print(numpy.size(a,1))
print(numpy.size(a,0))

print(numpy.prod([1, 10, 100, 5]))
//  [python]
def enable_pruning(self) -> None:
        """Enable pruning in the network. Pruning must be made permanent to recover pruned weights.

        Raises:
            ValueError: If the quantization parameters are invalid.
        """
        max_neuron_connections = self.max_active_neurons()

        if max_neuron_connections == 0:
            raise ValueError(
                "The maximum accumulator bit-width is too low "
                "for the quantization parameters requested. No neurons would be created in the "
                "requested configuration"
            )

        # Iterate over all layers that have weights (Linear ones)
        layer_idx = 0
        for layer in self.features:
            if not isinstance(layer, (nn.Linear, qnn.QuantLinear)):
                continue

            layer_shape = layer.weight.shape

            # Compute the fan-in, the number of inputs to a neuron, and the fan-out, the number of
            # neurons in the current layer.
            # The fan-in is the product of the kernel width x height x in_channels while the fan-out
            # is out_channels
            fan_in = n.prod(layer_shape[1:])
            fan_out = layer_shape[0]

            # To satisfy accumulator bit-width constraints each dot-product between an input line
            # and weight column must not exceed n_accum_bits bits. We thus prune the layer to have
            # at most max_neuron_connections non-zero weights
            if fan_in > max_neuron_connections and layer not in self.pruned_layers:
                pruning.l1_unstructured(
                    layer, "weight", (fan_in - max_neuron_connections) * fan_out
                )
                self.pruned_layers.add(layer)

            # If pruning is enabled, which is generally the case during training, a
            # forward hook is added in order to create a mask tensor (made of 0 or 1) that will
            # be multiplied with the weights during the forward pass.
            # This is done for all layers except the last one, which outputs the prediction
            if layer_idx < self.n_layers - 1 and self.n_prune_neurons_percentage > 0.0:
                # Use L2-norm structured pruning, using the torch ln_structured
                # function, with norm=2 and axis=0 (output/neuron axis)
                pruning.ln_structured(layer, "weight", self.n_prune_neurons_percentage, 2, 0)
                self.pruned_layers.add(layer)

            # Note this is counting only Linear layers
            layer_idx += 1

        assert_true(
            layer_idx == self.n_layers,
            "Not all layers in the network were examined as candidates for pruning",
        )
//  [python]
def enable_pruning(self) -> None:
        """Enable pruning in the network. Pruning must be made permanent to recover pruned weights.

        Raises:
            ValueError: If the quantization parameters are invalid.
        """
        max_neuron_connections = self.max_active_neurons()

        if max_neuron_connections == 0:
            raise ValueError(
                "The maximum accumulator bit-width is too low "
                "for the quantization parameters requested. No neurons would be created in the "
                "requested configuration"
            )

        # Iterate over all layers that have weights (Linear ones)
        layer_idx = 0
        for layer in self.features:
            if not isinstance(layer, (nn.Linear, qnn.QuantLinear)):
                continue

            layer_shape = layer.weight.shape

            # Compute the fan-in, the number of inputs to a neuron, and the fan-out, the number of
            # neurons in the current layer.
            # The fan-in is the product of the kernel width x height x in_channels while the fan-out
            # is out_channels
            fan_in = numpy.prod(layer_shape[1:])
            fan_out = layer_shape[0]

            # To satisfy accumulator bit-width constraints each dot-product between an input line
            # and weight column must not exceed n_accum_bits bits. We thus prune the layer to have
            # at most max_neuron_connections non-zero weights
            if fan_in > max_neuron_connections and layer not in self.pruned_layers:
                pruning.l1_unstructured(
                    layer, "weight", (fan_in - max_neuron_connections) * fan_out
                )
                self.pruned_layers.add(layer)

            # If pruning is enabled, which is generally the case during training, a
            # forward hook is added in order to create a mask tensor (made of 0 or 1) that will
            # be multiplied with the weights during the forward pass.
            # This is done for all layers except the last one, which outputs the prediction
            if layer_idx < self.n_layers - 1 and self.n_prune_neurons_percentage > 0.0:
                # Use L2-norm structured pruning, using the torch ln_structured
                # function, with norm=2 and axis=0 (output/neuron axis)
                pruning.ln_structured(layer, "weight", self.n_prune_neurons_percentage, 2, 0)
                self.pruned_layers.add(layer)

            # Note this is counting only Linear layers
            layer_idx += 1

        assert_true(
            layer_idx == self.n_layers,
            "Not all layers in the network were examined as candidates for pruning",
        )
//  [python]
def _set_to_folded(self):
        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
        fold_group_a = particle_grid_idx[:, :x_split].flatten()
        fold_group_b = np.flip(particle_grid_idx, axis=1)[:, :x_split].flatten()

        curr_pos = pyflex.get_positions().reshape((-1, 4))
        curr_pos[fold_group_a, :] = curr_pos[fold_group_b, :].copy()
        curr_pos[fold_group_a, 1] += 0.05  # group a particle position made tcurr_pos[self.fold_group_b, 1] + 0.05e at top of group b position.

        pyflex.set_positions(curr_pos)
        for i in range(10):
            pyflex.step()
        return self._get_info()['performance']
//  [python]
def __init__(self, shape: Union[int, Tuple[int]], dtype: type, name: str) -> None:
        self.shape = np.empty(shape).shape
        self.dtype = dtype
        self.name = name
        size = int(np.multiply.reduce(shape) * dtype(0).nbytes)
        self.shm = SharedMemory(name=name, size=size)
//  [python]
def __init__(self, shape: Union[int, Tuple[int]], dtype: type, name: str) -> None:
        self.shape = np.empty(shape).shape
        self.dtype = dtype
        self.name = name
        size = int(np.prod(shape) * dtype(0).nbytes)
        self.shm = SharedMemory(name=name, size=size)
//  [python]
def _reset(self):
        """ Right now only use one initial state. Need to make sure _reset always give the same result. Otherwise CEM will fail."""
        if hasattr(self, 'action_tool'):
            particle_pos = pyflex.get_positions().reshape(-1, 4)
            p1, p2, p3, p4 = self._get_key_point_idx()
            key_point_pos = particle_pos[(p1, p2), :3] # Was changed from from p1, p4.
            middle_point = np.mean(key_point_pos, axis=0)
            self.action_tool.reset([middle_point[0], 0.1, middle_point[2]])

            # picker_low = self.action_tool.picker_low
            # picker_high = self.action_tool.picker_high
            # offset_x = self.action_tool._get_pos()[0][0][0] - picker_low[0] - 0.3
            # picker_low[0] += offset_x
            # picker_high[0] += offset_x
            # picker_high[0] += 1.0
            # self.action_tool.update_picker_boundary(picker_low, picker_high)

        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
        # self.set_colors(colors) # TODO the phase actually changes the cloth dynamics so we do not change them for now. Maybe delete this later.

        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. Need to make sure _reset always give the same result. Otherwise CEM will fail."""
        if hasattr(self, 'action_tool'):
            particle_pos = pyflex.get_positions().reshape(-1, 4)
            p1, p2, p3, p4 = self._get_key_point_idx()
            key_point_pos = particle_pos[(p1, p2), :3] # Was changed from from p1, p4.
            middle_point = np.mean(key_point_pos, axis=0)
            self.action_tool.reset([middle_point[0], 0.1, middle_point[2]])

            # picker_low = self.action_tool.picker_low
            # picker_high = self.action_tool.picker_high
            # offset_x = self.action_tool._get_pos()[0][0][0] - picker_low[0] - 0.3
            # picker_low[0] += offset_x
            # picker_high[0] += offset_x
            # picker_high[0] += 1.0
            # self.action_tool.update_picker_boundary(picker_low, picker_high)

        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
        # self.set_colors(colors) # TODO the phase actually changes the cloth dynamics so we do not change them for now. Maybe delete this later.

        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 get_random_Nd_array_of_shape_and_type(test_Nd_array_shape, test_type=None):
    array_size = n.prod(test_Nd_array_shape)
    (lower, upper) = get_range_of_random_integers(test_type)
    arr = numpy.random.random_integers(lower, upper, array_size).reshape(test_Nd_array_shape)
    return arr if test_type is None else arr.astype(test_type)
//  [python]
def get_random_Nd_array_of_shape_and_type(test_Nd_array_shape, test_type=None):
    array_size = numpy.prod(test_Nd_array_shape)
    (lower, upper) = get_range_of_random_integers(test_type)
    arr = numpy.random.random_integers(lower, upper, array_size).reshape(test_Nd_array_shape)
    return arr if test_type is None else arr.astype(test_type)
//  [python]
def _build_random_sparse_tensor(shape, sparsity):
    array_size = numpy.prod(shape)
    num_sparse = int(array_size * sparsity)
    array = numpy.random.randn(num_sparse)
    indices = numpy.random.choice(
        numpy.arange(array_size), replace=False, size=num_sparse
    ).astype(numpy.int64)
    return SparseTensorProto(
        values=numpy_helper.from_array(array),
        indices=numpy_helper.from_array(indices),
        dims=shape,
    )
//  [python]
def rand_unique_int():
  randchoice = npr.RandomState(0).choice
  def fn(shape, dtype):
    return randchoice(onp.arange(onp.prod(shape), dtype=dtype),
                      size=shape, replace=False)
  return fn
//  [python]
def get_qkv_length(
            self,
            q_shape,
            k_shape,
            v_shape,
    ) -> Tuple[int]:
        """
        Args:
            q_shape(List[int]):
                q tensor shape.
            k_shape(List[int]):
                k tensor shape.
            v_shape(List[int]):
                v tensor shape.
        """
        q_n = n.prod(q_shape) + 1 if self.has_cls_embed else n.prod(q_shape)
        k_n = n.prod(k_shape) + 1 if self.has_cls_embed else n.prod(k_shape)
        v_n = n.prod(v_shape) + 1 if self.has_cls_embed else n.prod(v_shape)
        return q_n, k_n, v_n