Method split

caffe2/python/operator_test/segment_ops_test.py:172–186

(self, data, lengths, indices=None)

class LengthsTester(TesterBase):
    def split(self, data, lengths, indices=None):
        K = len(lengths)
        outputs = [
            np.zeros((lengths[seg_id], ) + data.shape[1:],
                     dtype=data.dtype) for seg_id in range(0, K)
        ]
        start = 0
        for i in range(0, K):
            for j in range(0, lengths[i]):
                data_index = start + j
                if indices is not None:
                    data_index = indices[data_index]
                outputs[i][j] = data[data_index]
            start += lengths[i]
        return outputs

    def unsplit(self, extra_shape, inputs, lengths):
        N = sum(lengths)