[Relax Frontend] An Pull Requests for Squeeze OpConverter

While deploying the bert model with relax frontend,I found a bug in Squeeze OpConvert,we could directly fetch the axis info if there are in attr argument. the PR is here: https://github.com/apache/tvm/pull/16059

and there is another strange problem while leaglize R.squeeze(),my model is bert-base-uncased-squad-v1.onnx, part of model is:

class Module:
    def main(input_ids: R.Tensor(("batch", "sequence"), dtype="int64"), attention_mask: R.Tensor(("batch", "sequence"), dtype="int64"), token_type_ids: R.Tensor(("batch", "sequence"), dtype="int64")) -> R.Tuple(R.Tensor(("batch", "sequence"), dtype="float32"), R.Tensor(("batch", "sequence"), dtype="float32")):
        batch = T.int64()
        sequence = T.int64()
        R.func_attr({"num_input": 3})
        cls = Module
        with R.dataflow():
            lv623: R.Tensor((batch, sequence, 1), dtype="float32") = lv622[0]
            lv624: R.Tensor((batch, sequence, 1), dtype="float32") = lv622[1]
            lv625: R.Tensor((batch, sequence), dtype="float32") = R.squeeze(lv623, axis=None)
            lv626: R.Tensor((batch, sequence), dtype="float32") = R.squeeze(lv624, axis=None)
            gv: R.Tuple(R.Tensor((batch, sequence), dtype="float32"), R.Tensor((batch, sequence), dtype="float32")) = lv625, lv626
        return gv

when we put axis=None to R.squeeze(), although the input tensor have symbol shape, the LegalizeOps() still dose not work, that means we could’t get correspondingly squeeze prim_func.

I have found some UT about R.squeeze() in test_transform_legalize_ops_manipulate.py, it works. so I doubt there may some thing wrong with LegalizeOps() if our give a quite complex IRModule.

def test_squeeze_no_axis():
# fmt: off
class Squeeze:
    def main(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) :
        gv: R.Tensor((2, 3, 1, 4), "float32") = R.squeeze(x)
        return gv

class Expected:
    def main(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) :
        gv = R.call_tir(Expected.squeeze, (x,), R.Tensor((2, 3, 4), dtype="float32"))
        return gv

    def squeeze(rxplaceholder: T.Buffer((T.int64(2), T.int64(1), T.int64(3), T.int64(1), T.int64(1), T.int64(4)), "float32"), T_squeeze: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32")):
        T.func_attr({"tir.noalias": True})
        for i0, i1, i2 in T.grid(T.int64(2), T.int64(3), T.int64(4)):
            with T.block("T_squeeze"):
                ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
                T.reads(rxplaceholder[ax0, T.int64(0), ax1, T.int64(0), T.int64(0), ax2])
                T.writes(T_squeeze[ax0, ax1, ax2])
                T_squeeze[ax0, ax1, ax2] = rxplaceholder[ax0, T.int64(0), ax1, T.int64(0), T.int64(0), ax2]
# fmt: on

mod = LegalizeOps()(Squeeze)
tvm.ir.assert_structural_equal(mod, Expected)

Please take a look if you have thoughts @MasterJH5574