Function test_split

tests/python/relax/test_frontend_tflite.py:222–250

()

def test_split():
    class Split(tf.Module):
        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
        def func(self, x):
            a, b, c = tf.split(x, 3, axis=1)
            return tf.raw_ops.Pack(values=[a, b, c], axis=1)

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((1, 30), dtype="float32")) -> R.Tensor((1, 3, 10), dtype="float32"):
            R.func_attr({"num_input": 1})
            with R.dataflow():
                lv: R.Tuple(
                    R.Tensor((1, 10), dtype="float32"),
                    R.Tensor((1, 10), dtype="float32"),
                    R.Tensor((1, 10), dtype="float32"),
                ) = R.split(x, indices_or_sections=3, axis=1)
                lv1: R.Tensor((1, 10), dtype="float32") = lv[0]
                lv2: R.Tensor((1, 1, 10), dtype="float32") = R.expand_dims(lv1, axis=[1])
                lv3: R.Tensor((1, 10), dtype="float32") = lv[1]
                lv4: R.Tensor((1, 1, 10), dtype="float32") = R.expand_dims(lv3, axis=[1])
                lv5: R.Tensor((1, 10), dtype="float32") = lv[2]
                lv6: R.Tensor((1, 1, 10), dtype="float32") = R.expand_dims(lv5, axis=[1])
                gv: R.Tensor((1, 3, 10), dtype="float32") = R.concat((lv2, lv4, lv6), axis=1)
                R.output(gv)
            return gv

    verify(Split, Expected)


def test_split_v_dynamic():