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.topk(
                sorted=False, flat=True
            )
            self.next = self.k
        # 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, numpy.prod(data.shape[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

        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.
        
util/runningstats.py