Function split

imperative/python/megengine/xla/rules/tensor.py:145–172

(inp, nsplit_or_sections, axis)

# if nsplit_or_sections is Sequence[int], means divide inputs into
# len(nsplit_or_sections) parts, and the i-th part has nsplit_or_sections[i] elements
def split(inp, nsplit_or_sections, axis):
    from .indexing import index_with_slices

    ishape = inp.shape
    if axis < 0:
        axis = axis + len(ishape)

    if isinstance(nsplit_or_sections, int):
        dimlen = ishape[axis]
        assert dimlen % nsplit_or_sections == 0, "not an equal division"
        sections = [dimlen // nsplit_or_sections] * nsplit_or_sections
    else:
        sections = nsplit_or_sections

    assert np.sum(sections) == ishape[axis], "error split param"

    slices = []
    start = 0
    for section in sections:
        slices.append(
            [
                None if idx != axis else slice(start, start + section, 1)
                for idx in range(len(ishape))
            ]
        )
        start = start + section

    return [index_with_slices(inp, slices[i]) for i in range(len(sections))]


@register_lower_rule(mops.Split, "Split")