Function test_mean

tests/python/nightly/test_nnapi/test_ops.py:262–288

()

def test_mean():
    remote_obj, tracker = remote()

    def create_model() -> tvm.IRModule:
        @tvm.script.ir_module
        class Module:
            @R.function
            def main(
                i0: R.Tensor((1, 10, 15), "float32"),
            ) -> R.Tensor((1, 10, 1), "float32"):
                n = T.int64()
                with R.dataflow():
                    t0: R.Tensor((1, 10, 1), "float32") = R.mean(i0, axis=[-1], keepdims=True)
                    R.output(t0)
                return t0

        return Module

    mod = create_model()
    verify(
        remote_obj,
        tracker,
        mod,
        inputs=[
            tvm.runtime.tensor(np.random.uniform(size=(1, 10, 15)).astype("float32")),
        ],
    )


def test_conv2d():