torch_backend.py 文件源码

def expand_dims(x, axis=-1):
    def _expand_dims(x, axis=axis):
        return torch.unsqueeze(x, axis)

    def _compute_output_shape(x, axis=axis):
        shape = list(_get_shape(x))
        shape.insert(axis, 1)
        return shape

    return get_op(_expand_dims, output_shape=_compute_output_shape, arguments=[axis])(x)