def METHOD_NAME():
    # Create a small relay module with multiple inputs to dispatch over.
    x = relay.var("x", shape=[10, 10], dtype="float32")
    w = relay.var("w", shape=[10, 10], dtype="float32")
    y = x + w
    mod = tvm.IRModule.from_expr(y)
    # Apply flexible dispatch to dim 1 for both inputs.
    mod = relay.transform.FlexibleShapeDispatch(axis=1, buckets=[5, 10], input_indices=[0, 1])(mod)
    # Compile and confirm that output shapes are correct.
    exe = relay.vm.compile(mod, "llvm")
    vm = runtime.vm.VirtualMachine(exe, tvm.cpu())
    x_w_5 = np.random.normal(size=[10, 5]).astype("float32")
    assert list(vm.invoke("main", x_w_5, x_w_5).shape) == [10, 5]
    x_w_10 = np.random.normal(size=[10, 10]).astype("float32")
    assert list(vm.invoke("main", x_w_10, x_w_10).shape) == [10, 10]