[GitHub] [tvm] juda opened a new pull request, #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

juda opened a new pull request, #12318:
URL: https://github.com/apache/tvm/pull/12318

   This PR provides two how-to guides to show the usage of
   
   1. `optimized_torch`: tuning a PyTorch model/function by MetaSchedule
   2. `as_torch`: wrap TVMscript into a PyTorch model/function


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] yelite commented on a diff in pull request #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

yelite commented on code in PR #12318:
URL: https://github.com/apache/tvm/pull/12318#discussion_r956890345


##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,172 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework which contains
+# optimized implementations of most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# One of the examples is to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 70,
+# the width is 80 and the kernel size is 20,
+# then the 1-d depthwise conv could be written in PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain Python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python call TVMscript,
+# which serves for an abstraction of program on various hardware backends.
+
+# As a concrete example, we can write such a TVMscript for 1-d depthwise conv code as below.
+# The computation procedure of `tvm_depthwise` is corresponding to the code snippet of `vanilla_depthwise`.
+
+# In our `tvm_depthwise` function, both inputs and outputs are set to be function parameters
+# that held on the multi-dimension buffers. For each buffer, the shape and data type information are required.
+# In the function body, there is a syntactic sugar `T.grid` for writing multiple nested iterators.
+# In the body of the loop, each computation is wrapped in an additional construct named `T.block`.
+# A block is a basic unit of computation. Inside the block, we need to provide a few more information about the block axes.
+# Here, 2 spatial and 1 reduce block iterators are created and bound to the loop iterators i, j and k.
+# The computations and machine learning compilation analysis will be defined around them.
+# The last 3 lines are computation statements, including an initialization of `C[vj, vi]` and the summing up along the axis k.
+# Finally, we place 2 decorators `T.prim_func` and `as_torch` above the definition of function,
+# which converts the Python AST to TVMscript AST and then converts to PyTorch's `nn.Module`.
+
+
+@as_torch
+@T.prim_func
+def tvm_depthwise(
+    A: T.Buffer((70, 80), "float32"),

Review Comment:
   Will it be useful to show a follow-up example of how to make shape a variable by having nested function?



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,172 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework which contains
+# optimized implementations of most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# One of the examples is to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 70,
+# the width is 80 and the kernel size is 20,
+# then the 1-d depthwise conv could be written in PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain Python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python call TVMscript,
+# which serves for an abstraction of program on various hardware backends.
+
+# As a concrete example, we can write such a TVMscript for 1-d depthwise conv code as below.
+# The computation procedure of `tvm_depthwise` is corresponding to the code snippet of `vanilla_depthwise`.
+
+# In our `tvm_depthwise` function, both inputs and outputs are set to be function parameters
+# that held on the multi-dimension buffers. For each buffer, the shape and data type information are required.
+# In the function body, there is a syntactic sugar `T.grid` for writing multiple nested iterators.
+# In the body of the loop, each computation is wrapped in an additional construct named `T.block`.
+# A block is a basic unit of computation. Inside the block, we need to provide a few more information about the block axes.
+# Here, 2 spatial and 1 reduce block iterators are created and bound to the loop iterators i, j and k.
+# The computations and machine learning compilation analysis will be defined around them.
+# The last 3 lines are computation statements, including an initialization of `C[vj, vi]` and the summing up along the axis k.
+# Finally, we place 2 decorators `T.prim_func` and `as_torch` above the definition of function,
+# which converts the Python AST to TVMscript AST and then converts to PyTorch's `nn.Module`.
+
+
+@as_torch
+@T.prim_func
+def tvm_depthwise(
+    A: T.Buffer((70, 80), "float32"),
+    B: T.Buffer((70, 20), "float32"),
+    C: T.Buffer((70, 61), "float32"),
+) -> None:
+    for j, i, k in T.grid(70, 61, 20):
+        with T.block():
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vj, vi] = T.float32(0)
+            C[vj, vi] += B[vj, vk] * A[vj, vi + vk]
+
+
+# We can build the TVMscript code by calling the `tune` method.
+# Without providing more information, the model will be tuned for CPU.
+
+tvm_depthwise.tune()

Review Comment:
   Will it be better to explicitly write down the default TuneConfig and target here so that reader has better idea on how to customize this?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] juda commented on pull request #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

juda commented on PR #12318:
URL: https://github.com/apache/tvm/pull/12318#issuecomment-1255006044

   Hi @masahi , I polished the tutorial according to your feedback, could you please read it one more time?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] juda commented on pull request #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

juda commented on PR #12318:
URL: https://github.com/apache/tvm/pull/12318#issuecomment-1240165651

   @tvm-bot rerun


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] masahi merged pull request #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

masahi merged PR #12318:
URL: https://github.com/apache/tvm/pull/12318


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] masahi commented on a diff in pull request #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

masahi commented on code in PR #12318:
URL: https://github.com/apache/tvm/pull/12318#discussion_r969252514


##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,171 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework which contains
+# optimized implementations of most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# For example, suppose that we are going to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 70,
+# the width is 80 and the kernel size is 20,
+# then the 1-d depthwise conv could be written in PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain Python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python called TVMscript,
+# serving for a high-level abstraction of TVM intermediate representative,

Review Comment:
   which serves as the high-level frontend for TVM's Tensor IR.



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,171 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework which contains
+# optimized implementations of most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# For example, suppose that we are going to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 70,
+# the width is 80 and the kernel size is 20,
+# then the 1-d depthwise conv could be written in PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain Python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python called TVMscript,
+# serving for a high-level abstraction of TVM intermediate representative,
+# which is easy to impose transformations and optimizations and deploy on various hardware backends.
+
+# As a concrete example, we can write such a TVMscript for 1-d depthwise conv code as below.
+# The computation procedure of `tvm_depthwise` is corresponding to the code snippet of `vanilla_depthwise`.

Review Comment:
   Replace the two sentences with:
   
   "The depthwise 1D convolution code above can be translated to TVMScript as follows."



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,176 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+To follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The PyTorch model to optimize, along with its example input, are provided by users.
+# We can optimized the PyTorch's module by calling the `optimized_torch` method in default setting.
+# Without providing extra information, the model will be tuned for CPU.
+
+simple_model = SimpleModel()
+example_input = torch.randn(20, 1, 10, 10)
+model_optimized_by_tvm = optimize_torch(simple_model, example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module.
+ret1 = model_optimized_by_tvm(example_input)
+
+torch.save(model_optimized_by_tvm, "model_optimized.pt")
+model_loaded = torch.load("model_optimized.pt")
+
+# We load the module and run it again.
+ret2 = model_loaded(example_input)
+
+# We show that the results from original SimpleModel,
+# optimized model and loaded model are the same.
+
+ret3 = simple_model.forward(example_input)
+testing.assert_allclose(ret1.detach().numpy(), ret2.detach().numpy(), atol=1e-5, rtol=1e-5)
+testing.assert_allclose(ret1.detach().numpy(), ret3.detach().numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Resnet18 optimized by TVM MetaSchedule
+# ------------------------------
+# In the following, we will show that our approach is able to
+# accelerate common and large models, such as Resnet18.
+
+# We will tune our model on the GPU.
+target_cuda = "nvidia/geforce-rtx-3070"
+
+# For PyTorch users, the nn.Module could be written as usual, except for
+# applying "optimize_torch" function on the resnet18 model.
+
+
+class MyResNet18(torch.nn.Module):
+    def __init__(self, target):
+        super(MyResNet18, self).__init__()
+        # Here we impose the `optimize_torch` function
+        # The default setting is adapted automatically by the number of operations of the optimized model.
+        self.resnet = optimize_torch(resnet18(), [torch.rand(1, 3, 224, 224)], target=target)
+
+    def forward(self, input):
+        return self.resnet(input)
+
+
+tvm_module_resnet18 = MyResNet18(target_cuda)
+
+
+######################################################################
+# Resnet18 optimized by TorchScript
+# ------------------------------
+# Let us write down a resnet18 model in a standard way.
+
+
+class JitModule(torch.nn.Module):
+    def __init__(self):
+        super(JitModule, self).__init__()
+        # Here we impose the `optimize_for_inference` function
+        # TorchScript also provides a built-in "optimize_for_inference" function to accelerate the inference.
+        self.resnet = torch.jit.optimize_for_inference(torch.jit.script(resnet18().cuda().eval()))
+
+    def forward(self, input):
+        return self.resnet(input)
+
+
+jit_module_resnet18 = JitModule()
+
+
+######################################################################
+# Compare the performance between two approaches.
+# ------------------------------
+# Using PyTorch's benchmark Compare class, we can have a direct comparison result between two inference models.

Review Comment:
   Drop this sentence.



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,171 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module

Review Comment:
   Replace all `TVMscript` in this file with `TVMScript`



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,171 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework which contains
+# optimized implementations of most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# For example, suppose that we are going to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 70,
+# the width is 80 and the kernel size is 20,
+# then the 1-d depthwise conv could be written in PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain Python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python called TVMscript,
+# serving for a high-level abstraction of TVM intermediate representative,
+# which is easy to impose transformations and optimizations and deploy on various hardware backends.

Review Comment:
   This sentence can be dropped



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,171 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_

Review Comment:
   No need to add me as an author.



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,171 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework which contains
+# optimized implementations of most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# For example, suppose that we are going to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 70,
+# the width is 80 and the kernel size is 20,
+# then the 1-d depthwise conv could be written in PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain Python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python called TVMscript,

Review Comment:
   in Python



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,171 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework which contains
+# optimized implementations of most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# For example, suppose that we are going to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 70,
+# the width is 80 and the kernel size is 20,
+# then the 1-d depthwise conv could be written in PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain Python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python called TVMscript,
+# serving for a high-level abstraction of TVM intermediate representative,
+# which is easy to impose transformations and optimizations and deploy on various hardware backends.
+
+# As a concrete example, we can write such a TVMscript for 1-d depthwise conv code as below.
+# The computation procedure of `tvm_depthwise` is corresponding to the code snippet of `vanilla_depthwise`.
+
+# In our `tvm_depthwise` function, both inputs and outputs are set to be function parameters
+# that are held on the multi-dimension buffers. For each buffer, the shape and data type information are required.
+# In the function body, the syntactic sugar `T.grid`, `T.block` and `T.axis.remap`
+# are for writing multiple nested iterators, wrapping a basic unit of computation, and
+# creating spatial or reduce block iterators and bounding to the loop iterators, respectively.
+# Inside the computation statement, `C[vj, vi]` is initialized and sums up along the axis k.
+# Finally, we place 2 decorators `T.prim_func` and `as_torch` above the definition of the function,
+# which converts the Python AST to TVMscript and then converts to PyTorch's `nn.Module`.

Review Comment:
   I think you only need to briefly explain `as_torch` decorator. Everything else in this paragraph can be dropped, since this is not a detailed tutorial for TVMScript.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,176 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+To follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The PyTorch model to optimize, along with its example input, are provided by users.
+# We can optimized the PyTorch's module by calling the `optimized_torch` method in default setting.

Review Comment:
   The PyTorch module will be tuned by TVM for the target hardware.



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,171 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework which contains
+# optimized implementations of most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# For example, suppose that we are going to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 70,
+# the width is 80 and the kernel size is 20,
+# then the 1-d depthwise conv could be written in PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain Python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python called TVMscript,
+# serving for a high-level abstraction of TVM intermediate representative,
+# which is easy to impose transformations and optimizations and deploy on various hardware backends.
+
+# As a concrete example, we can write such a TVMscript for 1-d depthwise conv code as below.
+# The computation procedure of `tvm_depthwise` is corresponding to the code snippet of `vanilla_depthwise`.
+
+# In our `tvm_depthwise` function, both inputs and outputs are set to be function parameters
+# that are held on the multi-dimension buffers. For each buffer, the shape and data type information are required.
+# In the function body, the syntactic sugar `T.grid`, `T.block` and `T.axis.remap`
+# are for writing multiple nested iterators, wrapping a basic unit of computation, and
+# creating spatial or reduce block iterators and bounding to the loop iterators, respectively.
+# Inside the computation statement, `C[vj, vi]` is initialized and sums up along the axis k.
+# Finally, we place 2 decorators `T.prim_func` and `as_torch` above the definition of the function,
+# which converts the Python AST to TVMscript and then converts to PyTorch's `nn.Module`.
+
+
+@as_torch
+@T.prim_func
+def tvm_depthwise(
+    A: T.Buffer((70, 80), "float32"),
+    B: T.Buffer((70, 20), "float32"),
+    C: T.Buffer((70, 61), "float32"),
+) -> None:
+    for j, i, k in T.grid(70, 61, 20):
+        with T.block():
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vj, vi] = T.float32(0)
+            C[vj, vi] += B[vj, vk] * A[vj, vi + vk]
+
+
+# We can build the TVMscript code by calling the `tune` method in default setting.
+# Without providing extra information, the model will be tuned for CPU.
+
+tvm_depthwise.tune()
+
+# We can print out the tuned TVMscript code to see how the program is transformed, as
+
+print(tvm_depthwise.script())
+
+# We can verify that the two outputs are the same:
+
+ret_tvm = torch.zeros(out_channel, width - kernel_size + 1)
+tvm_depthwise(inputs, filters, ret_tvm)
+
+testing.assert_allclose(ret_torch.cpu().numpy(), ret_tvm.cpu().numpy(), atol=1e-5, rtol=1e-5)
+
+
+######################################################################
+# Benchmark
+# -------------------------------
+# We will compare two operators by using PyTorch's benchmark toolkit.
+
+results = []
+for i in range(5):
+    inputs = torch.randn(out_channel, width)
+    filters = torch.randn(out_channel, kernel_size)
+    res = torch.zeros(out_channel, width - kernel_size + 1)
+    sub_label = f"[test {i}]"
+    results.append(
+        benchmark.Timer(
+            stmt="tvm_depthwise(inputs, filters, res)",
+            setup="from __main__ import tvm_depthwise",
+            globals={"inputs": inputs, "filters": filters, "res": res},
+            sub_label=sub_label,
+            description="TVMscript",
+        ).blocked_autorange()
+    )
+    results.append(
+        benchmark.Timer(
+            stmt="torch_depthwise(inputs, filters)",
+            setup="from __main__ import torch_depthwise",
+            globals={
+                "inputs": inputs,
+                "filters": filters,
+            },
+            sub_label=sub_label,
+            description="PyTorch",
+        ).blocked_autorange()
+    )
+compare = benchmark.Compare(results)
+compare.print()
+
+# In author's environment, the average inference time of `tvm_depthwise` is 120.0 us (TVM version is 0.9.0),

Review Comment:
   0.9.0 is the released version, I don't think this is the one you are using for development. There is no need to mention the TVM version.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,176 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+To follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.

Review Comment:
    We provide the `optimize_torch` function, which has a similar usage as `torch.jit.trace`.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,176 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_

Review Comment:
   Drop my name



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,176 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+To follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The PyTorch model to optimize, along with its example input, are provided by users.
+# We can optimized the PyTorch's module by calling the `optimized_torch` method in default setting.
+# Without providing extra information, the model will be tuned for CPU.
+
+simple_model = SimpleModel()
+example_input = torch.randn(20, 1, 10, 10)
+model_optimized_by_tvm = optimize_torch(simple_model, example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.

Review Comment:
   `nn.Module`



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,176 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+To follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule

Review Comment:
   Optimize



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,176 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+To follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The PyTorch model to optimize, along with its example input, are provided by users.
+# We can optimized the PyTorch's module by calling the `optimized_torch` method in default setting.
+# Without providing extra information, the model will be tuned for CPU.
+
+simple_model = SimpleModel()
+example_input = torch.randn(20, 1, 10, 10)
+model_optimized_by_tvm = optimize_torch(simple_model, example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module.
+ret1 = model_optimized_by_tvm(example_input)
+
+torch.save(model_optimized_by_tvm, "model_optimized.pt")
+model_loaded = torch.load("model_optimized.pt")
+
+# We load the module and run it again.
+ret2 = model_loaded(example_input)
+
+# We show that the results from original SimpleModel,
+# optimized model and loaded model are the same.
+
+ret3 = simple_model.forward(example_input)
+testing.assert_allclose(ret1.detach().numpy(), ret2.detach().numpy(), atol=1e-5, rtol=1e-5)
+testing.assert_allclose(ret1.detach().numpy(), ret3.detach().numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Resnet18 optimized by TVM MetaSchedule
+# ------------------------------
+# In the following, we will show that our approach is able to
+# accelerate common and large models, such as Resnet18.

Review Comment:
   Drop "large". Resnet18 is not a large model.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,176 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+To follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The PyTorch model to optimize, along with its example input, are provided by users.
+# We can optimized the PyTorch's module by calling the `optimized_torch` method in default setting.
+# Without providing extra information, the model will be tuned for CPU.
+
+simple_model = SimpleModel()
+example_input = torch.randn(20, 1, 10, 10)
+model_optimized_by_tvm = optimize_torch(simple_model, example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module.
+ret1 = model_optimized_by_tvm(example_input)
+
+torch.save(model_optimized_by_tvm, "model_optimized.pt")
+model_loaded = torch.load("model_optimized.pt")
+
+# We load the module and run it again.
+ret2 = model_loaded(example_input)
+
+# We show that the results from original SimpleModel,
+# optimized model and loaded model are the same.
+
+ret3 = simple_model.forward(example_input)
+testing.assert_allclose(ret1.detach().numpy(), ret2.detach().numpy(), atol=1e-5, rtol=1e-5)
+testing.assert_allclose(ret1.detach().numpy(), ret3.detach().numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Resnet18 optimized by TVM MetaSchedule
+# ------------------------------
+# In the following, we will show that our approach is able to
+# accelerate common and large models, such as Resnet18.
+
+# We will tune our model on the GPU.
+target_cuda = "nvidia/geforce-rtx-3070"
+
+# For PyTorch users, the nn.Module could be written as usual, except for
+# applying "optimize_torch" function on the resnet18 model.
+
+
+class MyResNet18(torch.nn.Module):
+    def __init__(self, target):
+        super(MyResNet18, self).__init__()
+        # Here we impose the `optimize_torch` function
+        # The default setting is adapted automatically by the number of operations of the optimized model.
+        self.resnet = optimize_torch(resnet18(), [torch.rand(1, 3, 224, 224)], target=target)
+
+    def forward(self, input):
+        return self.resnet(input)
+
+
+tvm_module_resnet18 = MyResNet18(target_cuda)
+
+
+######################################################################
+# Resnet18 optimized by TorchScript
+# ------------------------------
+# Let us write down a resnet18 model in a standard way.
+
+
+class JitModule(torch.nn.Module):
+    def __init__(self):
+        super(JitModule, self).__init__()
+        # Here we impose the `optimize_for_inference` function
+        # TorchScript also provides a built-in "optimize_for_inference" function to accelerate the inference.
+        self.resnet = torch.jit.optimize_for_inference(torch.jit.script(resnet18().cuda().eval()))
+
+    def forward(self, input):
+        return self.resnet(input)
+
+
+jit_module_resnet18 = JitModule()
+
+
+######################################################################
+# Compare the performance between two approaches.
+# ------------------------------
+# Using PyTorch's benchmark Compare class, we can have a direct comparison result between two inference models.
+
+results = []
+for i in range(5):
+    test_input = torch.rand(1, 3, 224, 224).cuda()
+    sub_label = f"[test {i}]"
+    results.append(
+        benchmark.Timer(
+            stmt="tvm_module_resnet18(test_input)",
+            setup="from __main__ import tvm_module_resnet18",
+            globals={"test_input": test_input},
+            sub_label=sub_label,
+            description="tuning by meta",
+        ).blocked_autorange()
+    )
+    results.append(
+        benchmark.Timer(
+            stmt="jit_module_resnet18(test_input)",
+            setup="from __main__ import jit_module_resnet18",
+            globals={"test_input": test_input},
+            sub_label=sub_label,
+            description="tuning by jit",
+        ).blocked_autorange()
+    )
+
+compare = benchmark.Compare(results)
+compare.print()
+
+# In author's environment, the average inference time of `tvm_module_resnet18` is 620.0 us (TVM version is 0.9.0),

Review Comment:
   Drop the reference to TVM version (see the same comment to `using_as_torch.py`)



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,176 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+To follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The PyTorch model to optimize, along with its example input, are provided by users.
+# We can optimized the PyTorch's module by calling the `optimized_torch` method in default setting.
+# Without providing extra information, the model will be tuned for CPU.
+
+simple_model = SimpleModel()
+example_input = torch.randn(20, 1, 10, 10)
+model_optimized_by_tvm = optimize_torch(simple_model, example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module.
+ret1 = model_optimized_by_tvm(example_input)
+
+torch.save(model_optimized_by_tvm, "model_optimized.pt")
+model_loaded = torch.load("model_optimized.pt")
+
+# We load the module and run it again.
+ret2 = model_loaded(example_input)
+
+# We show that the results from original SimpleModel,

Review Comment:
   "We show that the results from the original PyTorch model and the model tuned by TVM are the same, as expected."
   
   No need to verify the result of the "loaded model". Note that the term "loaded model" makes no sense. 



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,176 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+To follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The PyTorch model to optimize, along with its example input, are provided by users.
+# We can optimized the PyTorch's module by calling the `optimized_torch` method in default setting.
+# Without providing extra information, the model will be tuned for CPU.
+
+simple_model = SimpleModel()
+example_input = torch.randn(20, 1, 10, 10)
+model_optimized_by_tvm = optimize_torch(simple_model, example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module.
+ret1 = model_optimized_by_tvm(example_input)
+
+torch.save(model_optimized_by_tvm, "model_optimized.pt")
+model_loaded = torch.load("model_optimized.pt")
+
+# We load the module and run it again.
+ret2 = model_loaded(example_input)
+
+# We show that the results from original SimpleModel,
+# optimized model and loaded model are the same.
+
+ret3 = simple_model.forward(example_input)
+testing.assert_allclose(ret1.detach().numpy(), ret2.detach().numpy(), atol=1e-5, rtol=1e-5)
+testing.assert_allclose(ret1.detach().numpy(), ret3.detach().numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Resnet18 optimized by TVM MetaSchedule

Review Comment:
   Optimize Resnet18 by MetaSchedule



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,171 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework which contains
+# optimized implementations of most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# For example, suppose that we are going to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 70,
+# the width is 80 and the kernel size is 20,
+# then the 1-d depthwise conv could be written in PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain Python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python called TVMscript,
+# serving for a high-level abstraction of TVM intermediate representative,
+# which is easy to impose transformations and optimizations and deploy on various hardware backends.
+
+# As a concrete example, we can write such a TVMscript for 1-d depthwise conv code as below.
+# The computation procedure of `tvm_depthwise` is corresponding to the code snippet of `vanilla_depthwise`.
+
+# In our `tvm_depthwise` function, both inputs and outputs are set to be function parameters
+# that are held on the multi-dimension buffers. For each buffer, the shape and data type information are required.
+# In the function body, the syntactic sugar `T.grid`, `T.block` and `T.axis.remap`
+# are for writing multiple nested iterators, wrapping a basic unit of computation, and
+# creating spatial or reduce block iterators and bounding to the loop iterators, respectively.
+# Inside the computation statement, `C[vj, vi]` is initialized and sums up along the axis k.
+# Finally, we place 2 decorators `T.prim_func` and `as_torch` above the definition of the function,
+# which converts the Python AST to TVMscript and then converts to PyTorch's `nn.Module`.
+
+
+@as_torch
+@T.prim_func
+def tvm_depthwise(
+    A: T.Buffer((70, 80), "float32"),
+    B: T.Buffer((70, 20), "float32"),
+    C: T.Buffer((70, 61), "float32"),
+) -> None:
+    for j, i, k in T.grid(70, 61, 20):
+        with T.block():
+            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+            with T.init():
+                C[vj, vi] = T.float32(0)
+            C[vj, vi] += B[vj, vk] * A[vj, vi + vk]
+
+
+# We can build the TVMscript code by calling the `tune` method in default setting.
+# Without providing extra information, the model will be tuned for CPU.
+
+tvm_depthwise.tune()
+
+# We can print out the tuned TVMscript code to see how the program is transformed, as
+
+print(tvm_depthwise.script())
+
+# We can verify that the two outputs are the same:
+
+ret_tvm = torch.zeros(out_channel, width - kernel_size + 1)
+tvm_depthwise(inputs, filters, ret_tvm)
+
+testing.assert_allclose(ret_torch.cpu().numpy(), ret_tvm.cpu().numpy(), atol=1e-5, rtol=1e-5)
+
+
+######################################################################
+# Benchmark
+# -------------------------------
+# We will compare two operators by using PyTorch's benchmark toolkit.

Review Comment:
   Drop this sentence. It is not useful.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Define the resnet18 optimized by MetaSchedule
+# ------------------------------
+# In another example, we compare the two optimizers about the performance of resnet18
+# For learning how to define a resnet18 model via PyTorch's nn.Module,
+# you can refer to https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting
+
+# We will deploy our model on the GPU.
+# In the working machine, the GPU is nvidia/geforce-rtx-3070.
+target_cuda = "nvidia/geforce-rtx-3070"
+
+# The default setting is adapted automatically by the number of operations of the optimized model
+# When needed, we can define the configuration by ourselves, like:
+tuning_config = TuneConfig(
+    strategy="evolutionary",
+    num_trials_per_iter=64,
+    max_trials_per_task=20000,
+    max_trials_global=20000,
+)
+
+# For PyTorch users, the nn.Module could be written as usual, except for
+# applying "optimize_torch" function on the resnet18 model.
+# In such a way, we obtain a new resnet18 model optimized by MetaSchedule.
+
+
+class MyResNet18(torch.nn.Module):
+    def __init__(self, config, target=None):
+        super(MyResNet18, self).__init__()
+        self.means = torch.nn.Parameter(
+            torch.tensor([103.939, 116.779, 123.68]).resize_(1, 3, 1, 1)
+        ).cuda()
+        # Here we impose the `optimize_torch` function
+        self.resnet = optimize_torch(resnet18(), [torch.rand(1, 3, 224, 224)], config, target)
+
+    def forward(self, input):
+        return self.resnet(input - self.means)
+
+
+# Since we set the number of trials largely,
+# we might need to wait more time for the search.
+meta_module_resnet18 = MyResNet18(tuning_config, target_cuda)
+
+
+######################################################################
+# Define the resnet18 optimized by TorchScript
+# ------------------------------
+# Besides, let us define a resnet18 model in a standard way.
+# TorchScript also provides a built-in "optimize_for_inference" function to accelerate the inference,
+# we will compare the performance of those two optimizers later.
+
+
+class JitModule(torch.nn.Module):

Review Comment:
   Please address this comment. There is no need to have `JitModule`.
   



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,176 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+To follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The PyTorch model to optimize, along with its example input, are provided by users.
+# We can optimized the PyTorch's module by calling the `optimized_torch` method in default setting.
+# Without providing extra information, the model will be tuned for CPU.
+
+simple_model = SimpleModel()
+example_input = torch.randn(20, 1, 10, 10)
+model_optimized_by_tvm = optimize_torch(simple_model, example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module.
+ret1 = model_optimized_by_tvm(example_input)
+
+torch.save(model_optimized_by_tvm, "model_optimized.pt")
+model_loaded = torch.load("model_optimized.pt")
+
+# We load the module and run it again.
+ret2 = model_loaded(example_input)
+
+# We show that the results from original SimpleModel,
+# optimized model and loaded model are the same.
+
+ret3 = simple_model.forward(example_input)

Review Comment:
   Drop `forward`



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Define the resnet18 optimized by MetaSchedule
+# ------------------------------
+# In another example, we compare the two optimizers about the performance of resnet18
+# For learning how to define a resnet18 model via PyTorch's nn.Module,
+# you can refer to https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting
+
+# We will deploy our model on the GPU.
+# In the working machine, the GPU is nvidia/geforce-rtx-3070.
+target_cuda = "nvidia/geforce-rtx-3070"
+
+# The default setting is adapted automatically by the number of operations of the optimized model
+# When needed, we can define the configuration by ourselves, like:
+tuning_config = TuneConfig(
+    strategy="evolutionary",
+    num_trials_per_iter=64,
+    max_trials_per_task=20000,
+    max_trials_global=20000,
+)
+
+# For PyTorch users, the nn.Module could be written as usual, except for
+# applying "optimize_torch" function on the resnet18 model.
+# In such a way, we obtain a new resnet18 model optimized by MetaSchedule.
+
+
+class MyResNet18(torch.nn.Module):

Review Comment:
   Please address this comment. There is no need to have `MyResNet18`.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] juda commented on a diff in pull request #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

juda commented on code in PR #12318:
URL: https://github.com/apache/tvm/pull/12318#discussion_r939762972


##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,152 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your Tensor IR with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to wrap the Tensor IR code with PyTorch module.
+By the decorator `as_torch`, users are able to wrap a TVMscript code into an PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# while the performance in PyTorch is not fully satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,
+# and you need to define an 1-d depthwise convolution.
+# Assume the number of in_channel and out_channel are both 700,
+# the width is 800 and the kernel size is 50,
+# Then it could be written by PyTorch in one line:
+
+in_channel = 700
+out_channel = 700
+width = 800
+kernel_size = 50
+
+
+def torch_depthwise(inputs, filters):
+    return F.conv1d(inputs, filters.view(700, 1, 50), groups=700)
+
+
+# We can run this function as
+
+inputs = torch.randn(700, 800).cuda()
+filters = torch.randn(700, 50).cuda()
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain python code,
+# could be written as
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(700, 800 - 50 + 1).cuda()
+    for j in range(700):
+        for i in range(800 - 50 + 1):
+            for k in range(50):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# We plan to optimize the `depthwise` function by
+# leveraging the power of TVMscript.
+# A simple TVMscript code, can be rewritten as:
+
+
+@as_torch
+@tvm.script.ir_module
+class tvm_depthwise:
+    @T.prim_func
+    def f(a: T.handle, b: T.handle, c: T.handle) -> None:
+        T.func_attr({"global_symbol": "main", "tir.noalias": True})
+        A = T.match_buffer(a, (700, 800), "float32")
+        B = T.match_buffer(b, (700, 50), "float32")
+        C = T.match_buffer(c, (700, 751), "float32")
+        for j in T.thread_binding(0, 700, thread="blockIdx.x"):
+            for i in T.thread_binding(0, 751, thread="threadIdx.x"):
+                for k in range(50):
+                    with T.block("output"):
+                        C[j, i] += B[j, k] * A[j, i + k]
+
+
+# We can verify that two function are the same:
+
+ret_tvm = torch.zeros(700, 800 - 50 + 1).cuda()
+tvm_depthwise(inputs, filters, ret_tvm)
+
+testing.assert_allclose(ret_torch.cpu().numpy(), ret_tvm.cpu().numpy(), atol=1e-5, rtol=1e-5)
+
+# Tip: We also provide an optional method `tune(config, target)` for additional optimization.
+# In this case, users could call `tvm_depthwise.tune(target="nvidia/geforce-rtx-3070")`
+# for trying to tune the operators via TVM MetaSchedule.
+

Review Comment:
   In this example, calling `tvm_depthwise.tune(target="nvidia/geforce-rtx-3070")` doesn't bring a better performance, thus I mention it in a comment.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] juda commented on pull request #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

juda commented on PR #12318:
URL: https://github.com/apache/tvm/pull/12318#issuecomment-1233619332

   @masahi I have improved the text, could you please review it again?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] masahi commented on a diff in pull request #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

masahi commented on code in PR #12318:
URL: https://github.com/apache/tvm/pull/12318#discussion_r947625769


##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.

Review Comment:
   Nevertheless, sometimes you might want to write your own operators in PyTorch. In that case, the performance of such custom operators might not be satisfactory for your needs.



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.

Review Comment:
   PyTorch is a very popular machine learning framework which contains optimized implementations of most commonly used operators



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,
+# and you need to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 700,
+# the width is 800 and the kernel size is 50,
+# then the 1-d depthwise conv could be written by PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    global out_channel
+    global kernel_size
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# We plan to optimize the `depthwise` function by leveraging the power of TVMscript.
+# Firstly, we can write such a simple python code:
+
+
+@as_torch
+def tvm_depthwise_initializer(Channels: int, Width: int, Kernel: int, Output: int, dtype: str):
+    @T.prim_func
+    def f(a: T.handle, b: T.handle, c: T.handle) -> None:
+        T.func_attr({"global_symbol": "main", "tir.noalias": True})
+        A = T.match_buffer(a, (Channels, Width), dtype)
+        B = T.match_buffer(b, (Channels, Kernel), dtype)
+        C = T.match_buffer(c, (Channels, Output), dtype)
+        for j, i, k in T.grid(Channels, Output, Kernel):
+            with T.block():
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                with T.init():
+                    C[vj, vi] = T.float32(0)
+                C[vj, vi] += B[vj, vk] * A[vj, vi + vk]
+
+    return f
+
+
+# Then we fill out the parameters and generate the TVMscript program.
+
+tvm_depthwise = tvm_depthwise_initializer(
+    out_channel, width, kernel_size, width - kernel_size + 1, "float32"
+)
+
+# We can tune the TVMscript code by providing a target device.

Review Comment:
   There is no target device, this sentence doesn't make sense.



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,
+# and you need to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 700,
+# the width is 800 and the kernel size is 50,
+# then the 1-d depthwise conv could be written by PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    global out_channel
+    global kernel_size
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# We plan to optimize the `depthwise` function by leveraging the power of TVMscript.
+# Firstly, we can write such a simple python code:
+
+
+@as_torch
+def tvm_depthwise_initializer(Channels: int, Width: int, Kernel: int, Output: int, dtype: str):
+    @T.prim_func
+    def f(a: T.handle, b: T.handle, c: T.handle) -> None:
+        T.func_attr({"global_symbol": "main", "tir.noalias": True})
+        A = T.match_buffer(a, (Channels, Width), dtype)
+        B = T.match_buffer(b, (Channels, Kernel), dtype)
+        C = T.match_buffer(c, (Channels, Output), dtype)
+        for j, i, k in T.grid(Channels, Output, Kernel):
+            with T.block():
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                with T.init():
+                    C[vj, vi] = T.float32(0)
+                C[vj, vi] += B[vj, vk] * A[vj, vi + vk]
+
+    return f
+
+
+# Then we fill out the parameters and generate the TVMscript program.
+
+tvm_depthwise = tvm_depthwise_initializer(
+    out_channel, width, kernel_size, width - kernel_size + 1, "float32"
+)
+
+# We can tune the TVMscript code by providing a target device.
+# The model will deploy on CPU, and the optimization (e.g. tiling) will conduct automatically.
+
+tvm_depthwise.tune()
+
+# We can print out the tuned TVMscript code, as
+
+print(tvm_depthwise.script())
+
+# Hint: If user plan to deploy on GPU, the GPU target should be provided,

Review Comment:
   Again, there is no "target" in this tutorial. 



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,
+# and you need to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 700,
+# the width is 800 and the kernel size is 50,
+# then the 1-d depthwise conv could be written by PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    global out_channel
+    global kernel_size
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# We plan to optimize the `depthwise` function by leveraging the power of TVMscript.
+# Firstly, we can write such a simple python code:
+
+
+@as_torch
+def tvm_depthwise_initializer(Channels: int, Width: int, Kernel: int, Output: int, dtype: str):
+    @T.prim_func
+    def f(a: T.handle, b: T.handle, c: T.handle) -> None:
+        T.func_attr({"global_symbol": "main", "tir.noalias": True})
+        A = T.match_buffer(a, (Channels, Width), dtype)
+        B = T.match_buffer(b, (Channels, Kernel), dtype)
+        C = T.match_buffer(c, (Channels, Output), dtype)
+        for j, i, k in T.grid(Channels, Output, Kernel):
+            with T.block():
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                with T.init():
+                    C[vj, vi] = T.float32(0)
+                C[vj, vi] += B[vj, vk] * A[vj, vi + vk]
+
+    return f
+
+
+# Then we fill out the parameters and generate the TVMscript program.
+
+tvm_depthwise = tvm_depthwise_initializer(
+    out_channel, width, kernel_size, width - kernel_size + 1, "float32"
+)
+
+# We can tune the TVMscript code by providing a target device.
+# The model will deploy on CPU, and the optimization (e.g. tiling) will conduct automatically.
+
+tvm_depthwise.tune()
+
+# We can print out the tuned TVMscript code, as
+
+print(tvm_depthwise.script())
+
+# Hint: If user plan to deploy on GPU, the GPU target should be provided,
+# and all the PyTorch tensors should convert into GPU version.
+# The thread bindings will be added automatically.
+
+# We can verify that the two functions are the same:
+
+ret_tvm = torch.zeros(out_channel, width - kernel_size + 1)
+tvm_depthwise(inputs, filters, ret_tvm)
+
+testing.assert_allclose(ret_torch.cpu().numpy(), ret_tvm.cpu().numpy(), atol=1e-5, rtol=1e-5)
+
+
+######################################################################
+# Benchmark
+# -------------------------------
+# We will compare two operators by using PyTorch's benchmark toolkit.
+
+results = []
+for i in range(5):
+    inputs = torch.randn(out_channel, width)
+    filters = torch.randn(out_channel, kernel_size)
+    res = torch.zeros(out_channel, width - kernel_size + 1)
+    sub_label = f"[test {i}]"
+    results.append(
+        benchmark.Timer(
+            stmt="tvm_depthwise(inputs, filters, res)",
+            setup="from __main__ import tvm_depthwise",
+            globals={"inputs": inputs, "filters": filters, "res": res},
+            sub_label=sub_label,
+            description="TVMscript",
+        ).blocked_autorange()
+    )
+    results.append(
+        benchmark.Timer(
+            stmt="torch_depthwise(inputs, filters)",
+            setup="from __main__ import torch_depthwise",
+            globals={
+                "inputs": inputs,
+                "filters": filters,
+            },
+            sub_label=sub_label,
+            description="PyTorch",
+        ).blocked_autorange()
+    )
+compare = benchmark.Compare(results)
+compare.print()
+
+# In the working machine, the average inference time of `tvm_depthwise` is 120.0 us (TVM version is 0.9.0),

Review Comment:
   In author's environment, 



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,

Review Comment:
   It doesn't make a lot of sense to talk about a variant of Mobilenet (where only 2d convolution is used) but then suddenly bring up 1D convolution.



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,
+# and you need to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 700,
+# the width is 800 and the kernel size is 50,
+# then the 1-d depthwise conv could be written by PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    global out_channel
+    global kernel_size
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# We plan to optimize the `depthwise` function by leveraging the power of TVMscript.
+# Firstly, we can write such a simple python code:
+
+
+@as_torch
+def tvm_depthwise_initializer(Channels: int, Width: int, Kernel: int, Output: int, dtype: str):
+    @T.prim_func
+    def f(a: T.handle, b: T.handle, c: T.handle) -> None:
+        T.func_attr({"global_symbol": "main", "tir.noalias": True})
+        A = T.match_buffer(a, (Channels, Width), dtype)
+        B = T.match_buffer(b, (Channels, Kernel), dtype)
+        C = T.match_buffer(c, (Channels, Output), dtype)
+        for j, i, k in T.grid(Channels, Output, Kernel):
+            with T.block():
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                with T.init():
+                    C[vj, vi] = T.float32(0)
+                C[vj, vi] += B[vj, vk] * A[vj, vi + vk]
+
+    return f
+
+
+# Then we fill out the parameters and generate the TVMscript program.
+
+tvm_depthwise = tvm_depthwise_initializer(
+    out_channel, width, kernel_size, width - kernel_size + 1, "float32"
+)
+
+# We can tune the TVMscript code by providing a target device.
+# The model will deploy on CPU, and the optimization (e.g. tiling) will conduct automatically.
+
+tvm_depthwise.tune()
+
+# We can print out the tuned TVMscript code, as
+
+print(tvm_depthwise.script())
+
+# Hint: If user plan to deploy on GPU, the GPU target should be provided,
+# and all the PyTorch tensors should convert into GPU version.
+# The thread bindings will be added automatically.
+
+# We can verify that the two functions are the same:

Review Comment:
   two outputs



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,
+# and you need to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 700,
+# the width is 800 and the kernel size is 50,

Review Comment:
   the code uses kernel size 20



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,
+# and you need to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 700,
+# the width is 800 and the kernel size is 50,
+# then the 1-d depthwise conv could be written by PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    global out_channel
+    global kernel_size
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# We plan to optimize the `depthwise` function by leveraging the power of TVMscript.
+# Firstly, we can write such a simple python code:
+
+
+@as_torch
+def tvm_depthwise_initializer(Channels: int, Width: int, Kernel: int, Output: int, dtype: str):
+    @T.prim_func
+    def f(a: T.handle, b: T.handle, c: T.handle) -> None:
+        T.func_attr({"global_symbol": "main", "tir.noalias": True})
+        A = T.match_buffer(a, (Channels, Width), dtype)
+        B = T.match_buffer(b, (Channels, Kernel), dtype)
+        C = T.match_buffer(c, (Channels, Output), dtype)
+        for j, i, k in T.grid(Channels, Output, Kernel):
+            with T.block():
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                with T.init():
+                    C[vj, vi] = T.float32(0)
+                C[vj, vi] += B[vj, vk] * A[vj, vi + vk]
+
+    return f
+
+
+# Then we fill out the parameters and generate the TVMscript program.
+
+tvm_depthwise = tvm_depthwise_initializer(
+    out_channel, width, kernel_size, width - kernel_size + 1, "float32"
+)
+
+# We can tune the TVMscript code by providing a target device.
+# The model will deploy on CPU, and the optimization (e.g. tiling) will conduct automatically.

Review Comment:
   There are basic grammar issues here.  
   
   Instead of "deploy", just use "run". But more importantly, we want to say that the model will be "tuned" for CPU.



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,
+# and you need to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 700,
+# the width is 800 and the kernel size is 50,
+# then the 1-d depthwise conv could be written by PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    global out_channel
+    global kernel_size
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# We plan to optimize the `depthwise` function by leveraging the power of TVMscript.
+# Firstly, we can write such a simple python code:
+
+
+@as_torch
+def tvm_depthwise_initializer(Channels: int, Width: int, Kernel: int, Output: int, dtype: str):
+    @T.prim_func
+    def f(a: T.handle, b: T.handle, c: T.handle) -> None:
+        T.func_attr({"global_symbol": "main", "tir.noalias": True})
+        A = T.match_buffer(a, (Channels, Width), dtype)
+        B = T.match_buffer(b, (Channels, Kernel), dtype)
+        C = T.match_buffer(c, (Channels, Output), dtype)
+        for j, i, k in T.grid(Channels, Output, Kernel):
+            with T.block():
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                with T.init():
+                    C[vj, vi] = T.float32(0)
+                C[vj, vi] += B[vj, vk] * A[vj, vi + vk]
+
+    return f
+
+
+# Then we fill out the parameters and generate the TVMscript program.
+
+tvm_depthwise = tvm_depthwise_initializer(
+    out_channel, width, kernel_size, width - kernel_size + 1, "float32"
+)
+
+# We can tune the TVMscript code by providing a target device.
+# The model will deploy on CPU, and the optimization (e.g. tiling) will conduct automatically.
+
+tvm_depthwise.tune()
+
+# We can print out the tuned TVMscript code, as
+
+print(tvm_depthwise.script())
+
+# Hint: If user plan to deploy on GPU, the GPU target should be provided,
+# and all the PyTorch tensors should convert into GPU version.

Review Comment:
   This sentence alone doesn't make sense. 



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,
+# and you need to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 700,
+# the width is 800 and the kernel size is 50,
+# then the 1-d depthwise conv could be written by PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    global out_channel
+    global kernel_size

Review Comment:
   do you need `global` here?



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,
+# and you need to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 700,
+# the width is 800 and the kernel size is 50,
+# then the 1-d depthwise conv could be written by PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    global out_channel
+    global kernel_size
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# We plan to optimize the `depthwise` function by leveraging the power of TVMscript.
+# Firstly, we can write such a simple python code:
+
+
+@as_torch
+def tvm_depthwise_initializer(Channels: int, Width: int, Kernel: int, Output: int, dtype: str):
+    @T.prim_func
+    def f(a: T.handle, b: T.handle, c: T.handle) -> None:
+        T.func_attr({"global_symbol": "main", "tir.noalias": True})
+        A = T.match_buffer(a, (Channels, Width), dtype)
+        B = T.match_buffer(b, (Channels, Kernel), dtype)
+        C = T.match_buffer(c, (Channels, Output), dtype)
+        for j, i, k in T.grid(Channels, Output, Kernel):
+            with T.block():
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                with T.init():
+                    C[vj, vi] = T.float32(0)
+                C[vj, vi] += B[vj, vk] * A[vj, vi + vk]
+
+    return f
+
+
+# Then we fill out the parameters and generate the TVMscript program.
+
+tvm_depthwise = tvm_depthwise_initializer(
+    out_channel, width, kernel_size, width - kernel_size + 1, "float32"
+)
+
+# We can tune the TVMscript code by providing a target device.
+# The model will deploy on CPU, and the optimization (e.g. tiling) will conduct automatically.
+
+tvm_depthwise.tune()
+
+# We can print out the tuned TVMscript code, as
+
+print(tvm_depthwise.script())
+
+# Hint: If user plan to deploy on GPU, the GPU target should be provided,
+# and all the PyTorch tensors should convert into GPU version.
+# The thread bindings will be added automatically.

Review Comment:
   This line is not important, should be removed.



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,
+# and you need to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 700,
+# the width is 800 and the kernel size is 50,
+# then the 1-d depthwise conv could be written by PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    global out_channel
+    global kernel_size
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# We plan to optimize the `depthwise` function by leveraging the power of TVMscript.
+# Firstly, we can write such a simple python code:

Review Comment:
   I don't agree that this is simple python code. First of all, this is not python code, it is TVMScript. Since this is most likely the first encounter to TVMScript for new users, we should briefly explain what TVMScript is and the difference of this TVMScript code and truly simple code in `vanilla_depthwise`.
   
   And since `match_buffer` stuff is very noisy and not really relevant, we should use the syntax sugar like https://github.com/apache/tvm/blob/1a98ea2e4faf791eb8ec239209904ff96c016052/python/tvm/tir/tensor_intrin/x86.py#L29-L31 



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,
+# and you need to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 700,
+# the width is 800 and the kernel size is 50,
+# then the 1-d depthwise conv could be written by PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    global out_channel
+    global kernel_size
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# We plan to optimize the `depthwise` function by leveraging the power of TVMscript.
+# Firstly, we can write such a simple python code:
+
+
+@as_torch
+def tvm_depthwise_initializer(Channels: int, Width: int, Kernel: int, Output: int, dtype: str):
+    @T.prim_func
+    def f(a: T.handle, b: T.handle, c: T.handle) -> None:
+        T.func_attr({"global_symbol": "main", "tir.noalias": True})
+        A = T.match_buffer(a, (Channels, Width), dtype)
+        B = T.match_buffer(b, (Channels, Kernel), dtype)
+        C = T.match_buffer(c, (Channels, Output), dtype)
+        for j, i, k in T.grid(Channels, Output, Kernel):
+            with T.block():
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                with T.init():
+                    C[vj, vi] = T.float32(0)
+                C[vj, vi] += B[vj, vk] * A[vj, vi + vk]
+
+    return f
+
+
+# Then we fill out the parameters and generate the TVMscript program.
+
+tvm_depthwise = tvm_depthwise_initializer(
+    out_channel, width, kernel_size, width - kernel_size + 1, "float32"
+)
+
+# We can tune the TVMscript code by providing a target device.
+# The model will deploy on CPU, and the optimization (e.g. tiling) will conduct automatically.
+
+tvm_depthwise.tune()
+
+# We can print out the tuned TVMscript code, as
+
+print(tvm_depthwise.script())
+
+# Hint: If user plan to deploy on GPU, the GPU target should be provided,
+# and all the PyTorch tensors should convert into GPU version.
+# The thread bindings will be added automatically.
+
+# We can verify that the two functions are the same:
+
+ret_tvm = torch.zeros(out_channel, width - kernel_size + 1)
+tvm_depthwise(inputs, filters, ret_tvm)
+
+testing.assert_allclose(ret_torch.cpu().numpy(), ret_tvm.cpu().numpy(), atol=1e-5, rtol=1e-5)
+
+
+######################################################################
+# Benchmark
+# -------------------------------
+# We will compare two operators by using PyTorch's benchmark toolkit.
+
+results = []
+for i in range(5):
+    inputs = torch.randn(out_channel, width)
+    filters = torch.randn(out_channel, kernel_size)
+    res = torch.zeros(out_channel, width - kernel_size + 1)
+    sub_label = f"[test {i}]"
+    results.append(
+        benchmark.Timer(
+            stmt="tvm_depthwise(inputs, filters, res)",
+            setup="from __main__ import tvm_depthwise",
+            globals={"inputs": inputs, "filters": filters, "res": res},
+            sub_label=sub_label,
+            description="TVMscript",
+        ).blocked_autorange()
+    )
+    results.append(
+        benchmark.Timer(
+            stmt="torch_depthwise(inputs, filters)",
+            setup="from __main__ import torch_depthwise",
+            globals={
+                "inputs": inputs,
+                "filters": filters,
+            },
+            sub_label=sub_label,
+            description="PyTorch",
+        ).blocked_autorange()
+    )
+compare = benchmark.Compare(results)
+compare.print()
+
+# In the working machine, the average inference time of `tvm_depthwise` is 120.0 us (TVM version is 0.9.0),
+# while the average inference time of `torch_depthwise` is 210.0 us (PyTorch version is 1.11.0),
+# showing the performance arises by around 43%.

Review Comment:
   showing the speedup of around 43%



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework in which
+# it highly optimizes most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators
+# in PyTorch, but the performance could be not satisfactory.
+#
+# For example, assume you are writing a variance of MobileNet,
+# and you need to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 700,
+# the width is 800 and the kernel size is 50,
+# then the 1-d depthwise conv could be written by PyTorch in one line:

Review Comment:
   could be written in PyTorch



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] juda commented on a diff in pull request #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

juda commented on code in PR #12318:
URL: https://github.com/apache/tvm/pull/12318#discussion_r956898347


##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,172 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework which contains
+# optimized implementations of most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# One of the examples is to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 70,
+# the width is 80 and the kernel size is 20,
+# then the 1-d depthwise conv could be written in PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain Python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python call TVMscript,
+# which serves for an abstraction of program on various hardware backends.
+
+# As a concrete example, we can write such a TVMscript for 1-d depthwise conv code as below.
+# The computation procedure of `tvm_depthwise` is corresponding to the code snippet of `vanilla_depthwise`.
+
+# In our `tvm_depthwise` function, both inputs and outputs are set to be function parameters
+# that held on the multi-dimension buffers. For each buffer, the shape and data type information are required.
+# In the function body, there is a syntactic sugar `T.grid` for writing multiple nested iterators.
+# In the body of the loop, each computation is wrapped in an additional construct named `T.block`.
+# A block is a basic unit of computation. Inside the block, we need to provide a few more information about the block axes.
+# Here, 2 spatial and 1 reduce block iterators are created and bound to the loop iterators i, j and k.
+# The computations and machine learning compilation analysis will be defined around them.
+# The last 3 lines are computation statements, including an initialization of `C[vj, vi]` and the summing up along the axis k.
+# Finally, we place 2 decorators `T.prim_func` and `as_torch` above the definition of function,
+# which converts the Python AST to TVMscript AST and then converts to PyTorch's `nn.Module`.
+
+
+@as_torch
+@T.prim_func
+def tvm_depthwise(
+    A: T.Buffer((70, 80), "float32"),

Review Comment:
   If we pass shape variables then we need `match_buffer ` operators, which might confuse users. 
   Currently, I choose a minimal set of grammar.
   @masahi what's your idea?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] masahi commented on a diff in pull request #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

masahi commented on code in PR #12318:
URL: https://github.com/apache/tvm/pull/12318#discussion_r969284123


##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,176 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.

Review Comment:
   an introductory tutorial -> a tutorial



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,171 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module

Review Comment:
   Is "Wrap ... with " the right phrase? I think "Wrap ... as PyTorch Module" is more correct.



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,171 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.

Review Comment:
   This article is a tutorial on wrapping  TVMScript code as a PyTorch module.



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,171 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.

Review Comment:
   "Using the decorator..."
   
   Drop "a" before TVMScript
   
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] masahi commented on a diff in pull request #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

masahi commented on code in PR #12318:
URL: https://github.com/apache/tvm/pull/12318#discussion_r960153711


##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,172 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework which contains
+# optimized implementations of most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# One of the examples is to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 70,
+# the width is 80 and the kernel size is 20,
+# then the 1-d depthwise conv could be written in PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain Python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python call TVMscript,
+# which serves for an abstraction of program on various hardware backends.

Review Comment:
   I think calling TVMScript as "an abstraction of program on various hardware backends" is a bit long shot. I think it is a much more high-level, concrete thing.  



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,172 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework which contains
+# optimized implementations of most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# One of the examples is to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 70,
+# the width is 80 and the kernel size is 20,
+# then the 1-d depthwise conv could be written in PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain Python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python call TVMscript,
+# which serves for an abstraction of program on various hardware backends.
+
+# As a concrete example, we can write such a TVMscript for 1-d depthwise conv code as below.
+# The computation procedure of `tvm_depthwise` is corresponding to the code snippet of `vanilla_depthwise`.
+
+# In our `tvm_depthwise` function, both inputs and outputs are set to be function parameters
+# that held on the multi-dimension buffers. For each buffer, the shape and data type information are required.
+# In the function body, there is a syntactic sugar `T.grid` for writing multiple nested iterators.
+# In the body of the loop, each computation is wrapped in an additional construct named `T.block`.
+# A block is a basic unit of computation. Inside the block, we need to provide a few more information about the block axes.
+# Here, 2 spatial and 1 reduce block iterators are created and bound to the loop iterators i, j and k.
+# The computations and machine learning compilation analysis will be defined around them.
+# The last 3 lines are computation statements, including an initialization of `C[vj, vi]` and the summing up along the axis k.
+# Finally, we place 2 decorators `T.prim_func` and `as_torch` above the definition of function,
+# which converts the Python AST to TVMscript AST and then converts to PyTorch's `nn.Module`.

Review Comment:
   These sentences might be too detailed for a tutorial intended for PT users. I prefer a more succinct summary of what TVMScript is about, not necessarily explaining all the syntactic constructs used in the example.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)

Review Comment:
   meta would mean "Facebook" to most PT users.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Define the resnet18 optimized by MetaSchedule
+# ------------------------------
+# In another example, we compare the two optimizers about the performance of resnet18
+# For learning how to define a resnet18 model via PyTorch's nn.Module,
+# you can refer to https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting
+
+# We will deploy our model on the GPU.

Review Comment:
   deploy -> tune



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)

Review Comment:
   I think it's more important to compare result from (1) vanilla PT and (2) after conversion to TVM.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Define the resnet18 optimized by MetaSchedule
+# ------------------------------
+# In another example, we compare the two optimizers about the performance of resnet18
+# For learning how to define a resnet18 model via PyTorch's nn.Module,
+# you can refer to https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting
+
+# We will deploy our model on the GPU.
+# In the working machine, the GPU is nvidia/geforce-rtx-3070.
+target_cuda = "nvidia/geforce-rtx-3070"
+
+# The default setting is adapted automatically by the number of operations of the optimized model
+# When needed, we can define the configuration by ourselves, like:
+tuning_config = TuneConfig(
+    strategy="evolutionary",
+    num_trials_per_iter=64,
+    max_trials_per_task=20000,
+    max_trials_global=20000,
+)
+
+# For PyTorch users, the nn.Module could be written as usual, except for
+# applying "optimize_torch" function on the resnet18 model.
+# In such a way, we obtain a new resnet18 model optimized by MetaSchedule.

Review Comment:
   I don't know what you are trying to say here. How about just say: "To optimize the resnet18 model by MetaSchedule, you can apply ""optimize_torch" function.".



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Define the resnet18 optimized by MetaSchedule
+# ------------------------------
+# In another example, we compare the two optimizers about the performance of resnet18
+# For learning how to define a resnet18 model via PyTorch's nn.Module,
+# you can refer to https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting
+
+# We will deploy our model on the GPU.
+# In the working machine, the GPU is nvidia/geforce-rtx-3070.
+target_cuda = "nvidia/geforce-rtx-3070"
+
+# The default setting is adapted automatically by the number of operations of the optimized model
+# When needed, we can define the configuration by ourselves, like:
+tuning_config = TuneConfig(
+    strategy="evolutionary",
+    num_trials_per_iter=64,
+    max_trials_per_task=20000,
+    max_trials_global=20000,
+)
+
+# For PyTorch users, the nn.Module could be written as usual, except for
+# applying "optimize_torch" function on the resnet18 model.
+# In such a way, we obtain a new resnet18 model optimized by MetaSchedule.
+
+
+class MyResNet18(torch.nn.Module):
+    def __init__(self, config, target=None):
+        super(MyResNet18, self).__init__()
+        self.means = torch.nn.Parameter(
+            torch.tensor([103.939, 116.779, 123.68]).resize_(1, 3, 1, 1)
+        ).cuda()
+        # Here we impose the `optimize_torch` function
+        self.resnet = optimize_torch(resnet18(), [torch.rand(1, 3, 224, 224)], config, target)
+
+    def forward(self, input):
+        return self.resnet(input - self.means)
+
+
+# Since we set the number of trials largely,
+# we might need to wait more time for the search.
+meta_module_resnet18 = MyResNet18(tuning_config, target_cuda)
+
+
+######################################################################
+# Define the resnet18 optimized by TorchScript
+# ------------------------------
+# Besides, let us define a resnet18 model in a standard way.
+# TorchScript also provides a built-in "optimize_for_inference" function to accelerate the inference,
+# we will compare the performance of those two optimizers later.
+
+
+class JitModule(torch.nn.Module):
+    def __init__(self):
+        super(JitModule, self).__init__()
+        self.means = torch.nn.Parameter(
+            torch.tensor([103.939, 116.779, 123.68]).resize_(1, 3, 1, 1)
+        ).cuda()
+        # Here we impose the `optimize_for_inference` function
+        self.resnet = torch.jit.optimize_for_inference(torch.jit.script(resnet18().cuda().eval()))
+
+    def forward(self, input):
+        return self.resnet(input - self.means)
+
+
+jit_module_resnet18 = JitModule()
+
+######################################################################
+# Compare the performance between two scheduling approaches.
+# ------------------------------
+# Using PyTorch's benchmark Compare class, we can have a direct comparison result between two inference models.
+
+results = []
+for i in range(5):
+    test_input = torch.rand(1, 3, 224, 224).half().cuda()
+    sub_label = f"[test {i}]"
+    results.append(
+        benchmark.Timer(
+            stmt="meta_module_resnet18(test_input)",
+            setup="from __main__ import meta_module_resnet18",
+            globals={"test_input": test_input},
+            sub_label=sub_label,
+            description="tuning by meta",
+        ).blocked_autorange()
+    )
+    results.append(
+        benchmark.Timer(
+            stmt="jit_module_resnet18(test_input)",
+            setup="from __main__ import jit_module_resnet18",
+            globals={"test_input": test_input},
+            sub_label=sub_label,
+            description="tuning by jit",
+        ).blocked_autorange()
+    )
+
+# We can print the results on screen.
+compare = benchmark.Compare(results)
+compare.print()
+
+# In the working machine, the average inference time by `optimized_torch` is 860.5 us,
+# while the average inference time of `jit_optimized` is 1156.3 us,
+# showing the performance arises by around 1/4.

Review Comment:
   Apply my comment from `as_torch` tutorial here too. I won't repeat the same comment.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Define the resnet18 optimized by MetaSchedule
+# ------------------------------
+# In another example, we compare the two optimizers about the performance of resnet18
+# For learning how to define a resnet18 model via PyTorch's nn.Module,
+# you can refer to https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting
+
+# We will deploy our model on the GPU.
+# In the working machine, the GPU is nvidia/geforce-rtx-3070.

Review Comment:
   Drop this sentence. "working machine" is not a word in English. And "the GPU is nvidia/geforce-rtx-3070" doesn't add any new information.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Define the resnet18 optimized by MetaSchedule
+# ------------------------------
+# In another example, we compare the two optimizers about the performance of resnet18
+# For learning how to define a resnet18 model via PyTorch's nn.Module,
+# you can refer to https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting
+
+# We will deploy our model on the GPU.
+# In the working machine, the GPU is nvidia/geforce-rtx-3070.
+target_cuda = "nvidia/geforce-rtx-3070"
+
+# The default setting is adapted automatically by the number of operations of the optimized model
+# When needed, we can define the configuration by ourselves, like:
+tuning_config = TuneConfig(
+    strategy="evolutionary",
+    num_trials_per_iter=64,
+    max_trials_per_task=20000,
+    max_trials_global=20000,
+)
+
+# For PyTorch users, the nn.Module could be written as usual, except for
+# applying "optimize_torch" function on the resnet18 model.
+# In such a way, we obtain a new resnet18 model optimized by MetaSchedule.
+
+
+class MyResNet18(torch.nn.Module):

Review Comment:
   We shouldn't need `MyResnet18` boilerplate. Subtraction of `mean` from input is irrelevant for this tutorial, and this boilerplate just adds a noise.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Define the resnet18 optimized by MetaSchedule
+# ------------------------------
+# In another example, we compare the two optimizers about the performance of resnet18
+# For learning how to define a resnet18 model via PyTorch's nn.Module,
+# you can refer to https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting
+
+# We will deploy our model on the GPU.
+# In the working machine, the GPU is nvidia/geforce-rtx-3070.
+target_cuda = "nvidia/geforce-rtx-3070"
+
+# The default setting is adapted automatically by the number of operations of the optimized model
+# When needed, we can define the configuration by ourselves, like:

Review Comment:
   I understand what you mean here, but for new users who are not familiar with TVM or meta schedule, this sentence probably wouldn't make any sense. 



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.

Review Comment:
   The PyTorch model to optimize, along with its example input, are provided by users.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.

Review Comment:
   Drop "see the result". You are not showing "the result", whatever that is.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.

Review Comment:
   Explain what is "the default configuration" and what is the "first example".



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Define the resnet18 optimized by MetaSchedule
+# ------------------------------
+# In another example, we compare the two optimizers about the performance of resnet18
+# For learning how to define a resnet18 model via PyTorch's nn.Module,
+# you can refer to https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting
+
+# We will deploy our model on the GPU.
+# In the working machine, the GPU is nvidia/geforce-rtx-3070.
+target_cuda = "nvidia/geforce-rtx-3070"
+
+# The default setting is adapted automatically by the number of operations of the optimized model
+# When needed, we can define the configuration by ourselves, like:
+tuning_config = TuneConfig(
+    strategy="evolutionary",
+    num_trials_per_iter=64,
+    max_trials_per_task=20000,
+    max_trials_global=20000,
+)
+
+# For PyTorch users, the nn.Module could be written as usual, except for
+# applying "optimize_torch" function on the resnet18 model.
+# In such a way, we obtain a new resnet18 model optimized by MetaSchedule.
+
+
+class MyResNet18(torch.nn.Module):
+    def __init__(self, config, target=None):
+        super(MyResNet18, self).__init__()
+        self.means = torch.nn.Parameter(
+            torch.tensor([103.939, 116.779, 123.68]).resize_(1, 3, 1, 1)
+        ).cuda()
+        # Here we impose the `optimize_torch` function
+        self.resnet = optimize_torch(resnet18(), [torch.rand(1, 3, 224, 224)], config, target)
+
+    def forward(self, input):
+        return self.resnet(input - self.means)
+
+
+# Since we set the number of trials largely,
+# we might need to wait more time for the search.
+meta_module_resnet18 = MyResNet18(tuning_config, target_cuda)
+
+
+######################################################################
+# Define the resnet18 optimized by TorchScript
+# ------------------------------
+# Besides, let us define a resnet18 model in a standard way.
+# TorchScript also provides a built-in "optimize_for_inference" function to accelerate the inference,
+# we will compare the performance of those two optimizers later.

Review Comment:
   Same comment as "Define the resnet18 optimized by MetaSchedule" above. 



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Define the resnet18 optimized by MetaSchedule
+# ------------------------------
+# In another example, we compare the two optimizers about the performance of resnet18
+# For learning how to define a resnet18 model via PyTorch's nn.Module,
+# you can refer to https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting
+
+# We will deploy our model on the GPU.
+# In the working machine, the GPU is nvidia/geforce-rtx-3070.
+target_cuda = "nvidia/geforce-rtx-3070"
+
+# The default setting is adapted automatically by the number of operations of the optimized model
+# When needed, we can define the configuration by ourselves, like:
+tuning_config = TuneConfig(
+    strategy="evolutionary",
+    num_trials_per_iter=64,
+    max_trials_per_task=20000,
+    max_trials_global=20000,
+)
+
+# For PyTorch users, the nn.Module could be written as usual, except for
+# applying "optimize_torch" function on the resnet18 model.
+# In such a way, we obtain a new resnet18 model optimized by MetaSchedule.
+
+
+class MyResNet18(torch.nn.Module):
+    def __init__(self, config, target=None):
+        super(MyResNet18, self).__init__()
+        self.means = torch.nn.Parameter(
+            torch.tensor([103.939, 116.779, 123.68]).resize_(1, 3, 1, 1)
+        ).cuda()
+        # Here we impose the `optimize_torch` function
+        self.resnet = optimize_torch(resnet18(), [torch.rand(1, 3, 224, 224)], config, target)
+
+    def forward(self, input):
+        return self.resnet(input - self.means)
+
+
+# Since we set the number of trials largely,
+# we might need to wait more time for the search.

Review Comment:
   Please revise this sentence.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.

Review Comment:
   This is related to one of my comments earlier: Users who are new to `optimize_torch` don't know anything about `tuning_config` or `target`. So these two sentences are not useful.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Define the resnet18 optimized by MetaSchedule
+# ------------------------------
+# In another example, we compare the two optimizers about the performance of resnet18
+# For learning how to define a resnet18 model via PyTorch's nn.Module,
+# you can refer to https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting

Review Comment:
   Please revise these three sentences. Resnet18 is already defined by PT, we are not "defining" it. What are "the two optimizers"? Since this is a tutorial intended for PT users, we don't need to teach them how to define resnet18 in PT. And the link is irrelevant. 



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.

Review Comment:
   I think you copied "For us to follow this tutorial" from other tutorials, but this is not a good English phrase. We can just say "To follow this tutorial".



##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,172 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework which contains
+# optimized implementations of most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# One of the examples is to define a 1-d depthwise convolution operator.

Review Comment:
   "For example, suppose we want to define..."



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Define the resnet18 optimized by MetaSchedule
+# ------------------------------
+# In another example, we compare the two optimizers about the performance of resnet18
+# For learning how to define a resnet18 model via PyTorch's nn.Module,
+# you can refer to https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting
+
+# We will deploy our model on the GPU.
+# In the working machine, the GPU is nvidia/geforce-rtx-3070.
+target_cuda = "nvidia/geforce-rtx-3070"
+
+# The default setting is adapted automatically by the number of operations of the optimized model
+# When needed, we can define the configuration by ourselves, like:
+tuning_config = TuneConfig(
+    strategy="evolutionary",
+    num_trials_per_iter=64,
+    max_trials_per_task=20000,
+    max_trials_global=20000,
+)
+
+# For PyTorch users, the nn.Module could be written as usual, except for
+# applying "optimize_torch" function on the resnet18 model.
+# In such a way, we obtain a new resnet18 model optimized by MetaSchedule.
+
+
+class MyResNet18(torch.nn.Module):
+    def __init__(self, config, target=None):
+        super(MyResNet18, self).__init__()
+        self.means = torch.nn.Parameter(
+            torch.tensor([103.939, 116.779, 123.68]).resize_(1, 3, 1, 1)
+        ).cuda()
+        # Here we impose the `optimize_torch` function
+        self.resnet = optimize_torch(resnet18(), [torch.rand(1, 3, 224, 224)], config, target)
+
+    def forward(self, input):
+        return self.resnet(input - self.means)
+
+
+# Since we set the number of trials largely,
+# we might need to wait more time for the search.
+meta_module_resnet18 = MyResNet18(tuning_config, target_cuda)
+
+
+######################################################################
+# Define the resnet18 optimized by TorchScript
+# ------------------------------
+# Besides, let us define a resnet18 model in a standard way.
+# TorchScript also provides a built-in "optimize_for_inference" function to accelerate the inference,
+# we will compare the performance of those two optimizers later.
+
+
+class JitModule(torch.nn.Module):
+    def __init__(self):
+        super(JitModule, self).__init__()
+        self.means = torch.nn.Parameter(
+            torch.tensor([103.939, 116.779, 123.68]).resize_(1, 3, 1, 1)
+        ).cuda()
+        # Here we impose the `optimize_for_inference` function
+        self.resnet = torch.jit.optimize_for_inference(torch.jit.script(resnet18().cuda().eval()))
+
+    def forward(self, input):
+        return self.resnet(input - self.means)
+
+
+jit_module_resnet18 = JitModule()
+
+######################################################################
+# Compare the performance between two scheduling approaches.
+# ------------------------------
+# Using PyTorch's benchmark Compare class, we can have a direct comparison result between two inference models.
+
+results = []
+for i in range(5):
+    test_input = torch.rand(1, 3, 224, 224).half().cuda()
+    sub_label = f"[test {i}]"
+    results.append(
+        benchmark.Timer(
+            stmt="meta_module_resnet18(test_input)",
+            setup="from __main__ import meta_module_resnet18",
+            globals={"test_input": test_input},
+            sub_label=sub_label,
+            description="tuning by meta",
+        ).blocked_autorange()
+    )
+    results.append(
+        benchmark.Timer(
+            stmt="jit_module_resnet18(test_input)",
+            setup="from __main__ import jit_module_resnet18",
+            globals={"test_input": test_input},
+            sub_label=sub_label,
+            description="tuning by jit",
+        ).blocked_autorange()
+    )
+
+# We can print the results on screen.

Review Comment:
   Drop this sentence



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Define the resnet18 optimized by MetaSchedule
+# ------------------------------
+# In another example, we compare the two optimizers about the performance of resnet18
+# For learning how to define a resnet18 model via PyTorch's nn.Module,
+# you can refer to https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting
+
+# We will deploy our model on the GPU.
+# In the working machine, the GPU is nvidia/geforce-rtx-3070.
+target_cuda = "nvidia/geforce-rtx-3070"
+
+# The default setting is adapted automatically by the number of operations of the optimized model
+# When needed, we can define the configuration by ourselves, like:
+tuning_config = TuneConfig(
+    strategy="evolutionary",
+    num_trials_per_iter=64,
+    max_trials_per_task=20000,
+    max_trials_global=20000,
+)
+
+# For PyTorch users, the nn.Module could be written as usual, except for
+# applying "optimize_torch" function on the resnet18 model.
+# In such a way, we obtain a new resnet18 model optimized by MetaSchedule.
+
+
+class MyResNet18(torch.nn.Module):
+    def __init__(self, config, target=None):
+        super(MyResNet18, self).__init__()
+        self.means = torch.nn.Parameter(
+            torch.tensor([103.939, 116.779, 123.68]).resize_(1, 3, 1, 1)
+        ).cuda()
+        # Here we impose the `optimize_torch` function
+        self.resnet = optimize_torch(resnet18(), [torch.rand(1, 3, 224, 224)], config, target)
+
+    def forward(self, input):
+        return self.resnet(input - self.means)
+
+
+# Since we set the number of trials largely,
+# we might need to wait more time for the search.
+meta_module_resnet18 = MyResNet18(tuning_config, target_cuda)
+
+
+######################################################################
+# Define the resnet18 optimized by TorchScript
+# ------------------------------
+# Besides, let us define a resnet18 model in a standard way.
+# TorchScript also provides a built-in "optimize_for_inference" function to accelerate the inference,
+# we will compare the performance of those two optimizers later.
+
+
+class JitModule(torch.nn.Module):

Review Comment:
   Drop `JitModule` boilerplate.



##########
gallery/how_to/work_with_pytorch/using_optimized_torch.py:
##########
@@ -0,0 +1,193 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Compile PyTorch Models
+======================
+**Author**: `Yaoda Zhou <https://github.com/juda/>`_
+This article is an introductory tutorial to optimize PyTorch models by using `tvm.contrib.torch.optimize_torch`.
+For us to follow this tutorial, PyTorch, as well as TorchVision, should be installed.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Import library for profiling
+import torch.utils.benchmark as benchmark
+from torchvision.models import resnet18
+
+# Import `optimize_torch` function
+from tvm.contrib.torch import optimize_torch
+from tvm.meta_schedule import TuneConfig
+
+######################################################################
+# Define a simple module written by PyTorch
+# ------------------------------
+
+
+class SimpleModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5)
+        self.conv2 = nn.Conv2d(20, 20, 5)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        return F.relu(self.conv2(x))
+
+
+######################################################################
+# Optimized SimpleModel by TVM MetaSchedule
+# ------------------------------
+# We provide a `optimize_torch` function, which has the similar usage as `torch.jit.trace`.
+# The optimized function/model and example input are required to provide by users.
+# If the third parameter `tuning_config` is not provided, a default configuration is loaded.
+# If the parameter `target` is empty, the model will deploy on the CPU.
+
+
+example_input = torch.randn(20, 1, 10, 10)
+
+# We use the default configuration for the first example.
+model_optimized_by_meta = optimize_torch(SimpleModel(), example_input)
+
+######################################################################
+# Save/Load module
+# ------------------------------
+# We can save and load our tuned module like the standard `nn.module`.
+
+# Let us run our tuned module and see the result.
+ret1 = model_optimized_by_meta(example_input)
+
+torch.save(model_optimized_by_meta, "meta_model.pt")
+model_loaded = torch.load("meta_model.pt")
+
+# We load the module and run it again, and it will return the same result as above.
+ret2 = model_loaded(example_input)
+
+testing.assert_allclose(ret1.numpy(), ret2.numpy(), atol=1e-5, rtol=1e-5)
+
+######################################################################
+# Define the resnet18 optimized by MetaSchedule
+# ------------------------------
+# In another example, we compare the two optimizers about the performance of resnet18
+# For learning how to define a resnet18 model via PyTorch's nn.Module,
+# you can refer to https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting
+
+# We will deploy our model on the GPU.
+# In the working machine, the GPU is nvidia/geforce-rtx-3070.
+target_cuda = "nvidia/geforce-rtx-3070"
+
+# The default setting is adapted automatically by the number of operations of the optimized model
+# When needed, we can define the configuration by ourselves, like:
+tuning_config = TuneConfig(
+    strategy="evolutionary",
+    num_trials_per_iter=64,
+    max_trials_per_task=20000,
+    max_trials_global=20000,
+)
+
+# For PyTorch users, the nn.Module could be written as usual, except for
+# applying "optimize_torch" function on the resnet18 model.
+# In such a way, we obtain a new resnet18 model optimized by MetaSchedule.
+
+
+class MyResNet18(torch.nn.Module):
+    def __init__(self, config, target=None):
+        super(MyResNet18, self).__init__()
+        self.means = torch.nn.Parameter(
+            torch.tensor([103.939, 116.779, 123.68]).resize_(1, 3, 1, 1)
+        ).cuda()
+        # Here we impose the `optimize_torch` function
+        self.resnet = optimize_torch(resnet18(), [torch.rand(1, 3, 224, 224)], config, target)
+
+    def forward(self, input):
+        return self.resnet(input - self.means)
+
+
+# Since we set the number of trials largely,
+# we might need to wait more time for the search.
+meta_module_resnet18 = MyResNet18(tuning_config, target_cuda)
+
+
+######################################################################
+# Define the resnet18 optimized by TorchScript
+# ------------------------------
+# Besides, let us define a resnet18 model in a standard way.
+# TorchScript also provides a built-in "optimize_for_inference" function to accelerate the inference,
+# we will compare the performance of those two optimizers later.
+
+
+class JitModule(torch.nn.Module):
+    def __init__(self):
+        super(JitModule, self).__init__()
+        self.means = torch.nn.Parameter(
+            torch.tensor([103.939, 116.779, 123.68]).resize_(1, 3, 1, 1)
+        ).cuda()
+        # Here we impose the `optimize_for_inference` function
+        self.resnet = torch.jit.optimize_for_inference(torch.jit.script(resnet18().cuda().eval()))
+
+    def forward(self, input):
+        return self.resnet(input - self.means)
+
+
+jit_module_resnet18 = JitModule()
+
+######################################################################
+# Compare the performance between two scheduling approaches.

Review Comment:
   What are "two scheduling approaches"? `torch.jit.optimize_for_inference` is not a scheduling approach.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] masahi commented on a diff in pull request #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

masahi commented on code in PR #12318:
URL: https://github.com/apache/tvm/pull/12318#discussion_r956933885


##########
gallery/how_to/work_with_pytorch/using_as_torch.py:
##########
@@ -0,0 +1,172 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Wrap Your TVMscript with PyTorch Module
+======================
+**Author**: 
+`Yaoda Zhou <https://github.com/juda>`_,
+`Masahiro Masuda <https://github.com/masahi>`_
+
+This article is an introductory tutorial on wrapping the TVMscript code with the PyTorch module.
+By the decorator `as_torch`, users can wrap a TVMscript code into a PyTorch nn.Module naturally.
+"""
+
+# sphinx_gallery_start_ignore
+from tvm import testing
+
+testing.utils.install_request_hook(depth=3)
+# sphinx_gallery_end_ignore
+
+# Import PyTorch, as well as necessary libraries
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+
+import tvm
+from tvm.contrib.torch import as_torch
+from tvm.script import tir as T
+
+######################################################################
+# Write your own PyTorch operator by TVMscript
+# -------------------------------
+# PyTorch is a very popular machine learning framework which contains
+# optimized implementations of most commonly used operators.
+# Nevertheless, sometimes you might want to write your own operators in PyTorch.
+# In that case, the performance of such custom operators might not be satisfactory for your needs.
+#
+# One of the examples is to define a 1-d depthwise convolution operator.
+# Assume the number of in_channel and out_channel are both 70,
+# the width is 80 and the kernel size is 20,
+# then the 1-d depthwise conv could be written in PyTorch in one line:
+
+in_channel = 70
+out_channel = 70
+width = 80
+kernel_size = 20
+
+
+def torch_depthwise(inputs, filters):
+    return F.conv1d(inputs, filters.view(out_channel, 1, kernel_size), groups=out_channel)
+
+
+# We can run this function as:
+
+inputs = torch.randn(in_channel, width)
+filters = torch.randn(out_channel, kernel_size)
+ret_torch = torch_depthwise(inputs, filters)
+
+# The `torch_depthwise` function, in a plain Python code, could be written as:
+
+
+def vanilla_depthwise(input, weight):
+    ret = torch.zeros(out_channel, width - kernel_size + 1)
+    for j in range(out_channel):
+        for i in range(width - kernel_size + 1):
+            for k in range(kernel_size):
+                ret[j, i] += weight[j, k] * input[j, i + k]
+    return ret
+
+
+# Then, we plan to optimize the `depthwise` function by leveraging the power of TVM.
+# TVM community proposes an embedded Domain Specific Language on Python call TVMscript,
+# which serves for an abstraction of program on various hardware backends.
+
+# As a concrete example, we can write such a TVMscript for 1-d depthwise conv code as below.
+# The computation procedure of `tvm_depthwise` is corresponding to the code snippet of `vanilla_depthwise`.
+
+# In our `tvm_depthwise` function, both inputs and outputs are set to be function parameters
+# that held on the multi-dimension buffers. For each buffer, the shape and data type information are required.
+# In the function body, there is a syntactic sugar `T.grid` for writing multiple nested iterators.
+# In the body of the loop, each computation is wrapped in an additional construct named `T.block`.
+# A block is a basic unit of computation. Inside the block, we need to provide a few more information about the block axes.
+# Here, 2 spatial and 1 reduce block iterators are created and bound to the loop iterators i, j and k.
+# The computations and machine learning compilation analysis will be defined around them.
+# The last 3 lines are computation statements, including an initialization of `C[vj, vi]` and the summing up along the axis k.
+# Finally, we place 2 decorators `T.prim_func` and `as_torch` above the definition of function,
+# which converts the Python AST to TVMscript AST and then converts to PyTorch's `nn.Module`.
+
+
+@as_torch
+@T.prim_func
+def tvm_depthwise(
+    A: T.Buffer((70, 80), "float32"),

Review Comment:
   I think the `Buffer` syntax sugar can be extended for dynamic shapes. But currently we cannot tune over dynamic shapes, so the performance will probably be slower than PT.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] masahi commented on pull request #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

masahi commented on PR #12318:
URL: https://github.com/apache/tvm/pull/12318#issuecomment-1259033464

   @juda Thank your for your patience. I think the tutorial is now very clean and simple, without unnecessary things.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [tvm] juda commented on pull request #12318: [TVM PyTorch Integration] optimized_torch & as_torch how-to guide

[GitBox <gi...@apache.org>]

juda commented on PR #12318:
URL: https://github.com/apache/tvm/pull/12318#issuecomment-1241582430

   @masahi I finish another round of polishing. Could you please have a look?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org