Function stack

imperative/python/megengine/xla/rules/tensor.py:126–129

(inps, axis)

def stack(inps, axis):
    assert len(inps) > 0, f"stack inputs should not be empty"
    inps = [expand_dims(inp, axis) for inp in inps]
    return concat(inps, axis)


@register_lower_rule(mops.Concat, "Concat")