

def get_partitions(self, x, slices):
    '\n        given a partitioning, slices, of an activation buffer, x, determine which axis to slice\n        along depending on whether x is a sequential tensor or not\n        '
    if (x.shape[(- 1)] != self.be.bsz):
        return [x[:, sl] for sl in slices]
    else:
        return [x[sl] for sl in slices]

dataset/ETHPy150Open NervanaSystems/neon/neon/layers/container.py MergeBroadcast.get_partitions/original