sunggg opened a new pull request, #15221:
URL: https://github.com/apache/tvm/pull/15221
This PR introduces an utility pass that binds symbolic variables to
user-provided integer values.
For example, say we have a following IRModule.
```Python
@tvm.script.ir_module
class Before:
@R.function
def main(
x: R.Tensor(("m", "n")),
y: R.Tensor(("m", "n"))
) -> R.Tensor(("m", "n")):
m = T.Var("m", "int64")
n = T.Var("m", "int64")
with R.dataflow():
out = R.matmul(x, y)
R.output(out)
return out
```
We can conveniently bind the symbolic variable by applying `After =
relax.transform.BindSymVars("main", {"m": 10, "n": 10})(Before)`.
```Python
@tvm.script.ir_module
class After:
@R.function
def main(
x: R.Tensor((10, 10)),
y: R.Tensor((10, 10))
) -> R.Tensor((10, 10)):
with R.dataflow():
out = R.matmul(x, y)
R.output(out)
return out
```
This would be useful when providing compile-time shape info (e.g., model
params or batch sizes) by eliminating the need to rewrite the model.
cc. @tqchen @psrivas2
--
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]
