Function expand_dims

imperative/python/megengine/xla/rules/tensor.py:81–91

(inp, axis)

def expand_dims(inp, axis):
    assert isinstance(axis, int), f"only int axis supported, get {axis}"
    assert (
        axis >= -inp.ndim - 1 and axis <= inp.ndim
    ), f"invalid axis {axis} for {inp.shape}"

    dst_shape = list(inp.shape)
    insert_pos = axis if axis >= 0 else (axis + inp.ndim + 1)
    dst_shape.insert(insert_pos, 1)

    return inp.reshape(tuple(dst_shape))


@register_lower_rule(mops.Dimshuffle)