yongwww commented on PR #14123:
URL: https://github.com/apache/tvm/pull/14123#issuecomment-1446683807

   
   > Thanks @yongwww , as a followup, would be great if we can also confirm if 
it works on TVMScript parser
   
   Thanks @tqchen! Is the following test case what you meant?
   
   ```python
   @tvm.script.ir_module
       class Before:
           @R.function
           def main(x: R.Tensor((10, 20), "float32")) -> R.Tensor(None, 
"float32", ndim=2):
               # block 0
               with R.dataflow():
                   lv0 = R.emit_te(topi.add, x, relax.const(1, "float32"))
                   lv1 = R.emit_te(topi.exp, lv0)
                   gv = R.emit_te(topi.squeeze, lv1)
                   R.output(gv)
               return gv
   
       Before.show()
   ```
   
   output of Before.show():
   ```python
   @I.ir_module
   class Module:
       @R.function
       def main(
           x: R.Tensor((10, 20), dtype="float32")
       ) -> R.Tensor((10, 20), dtype="float32"):
           with R.dataflow():
               lv = R.call_tir(
                   add,
                   (x, R.const(1, "float32")),
                   out_sinfo=R.Tensor((10, 20), dtype="float32"),
               )
               lv0: R.Tensor((10, 20), dtype="float32") = lv
               lv2 = R.call_tir(exp, (lv0,), out_sinfo=R.Tensor((10, 20), 
dtype="float32"))
               lv1: R.Tensor((10, 20), dtype="float32") = lv2
               lv4 = R.call_tir(
                   squeeze, (lv1,), out_sinfo=R.Tensor((10, 20), 
dtype="float32")
               )
               gv: R.Tensor((10, 20), dtype="float32") = lv4
               R.output(gv)
           return gv
   ```
   
   
   


-- 
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: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to