sunggg opened a new pull request, #15246:
URL: https://github.com/apache/tvm/pull/15246
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 specializing shape and 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]