This is an automated email from the ASF dual-hosted git repository.
sanirudh pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 0328e98b94 [Unity][NN] Allow nn.Placeholder/Parameter prior to
BlockBuilder (#15025)
0328e98b94 is described below
commit 0328e98b940adb970185ce97f015e924b640ba09
Author: Eric Lunderberg <[email protected]>
AuthorDate: Fri Jun 9 01:03:15 2023 -0400
[Unity][NN] Allow nn.Placeholder/Parameter prior to BlockBuilder (#15025)
* [Unity][NN] Allow nn.Placeholder/Parameter prior to BlockBuilder
Prior to this commit, use of `nn.Placeholder` or `nn.Parameter`
outside of a `with block_builder.function('name'):` scope resulted in
an error. This commit updates the behavior to allow declaration prior
to entering the `with` block. This can be useful for declaring a
model object, which is then used to define several related functions.
The scope was required so that `relax.BlockBuilder.current()` could
de-duplicate variable names. While two distinct variables in Relax
may have identical names, for user readability it is convenient to
have all names be unique within a Relax function. This commit
maintains the de-duplication of names if a `nn.Placeholder` or
`nn.Parameter` is defined within an active `relax.BlockBuilder`, that
context may be used to provide a unique name.
* Lint fix
---
python/tvm/relax/testing/nn.py | 43 ++++++++++++++++++++++++++++++++++++------
1 file changed, 37 insertions(+), 6 deletions(-)
diff --git a/python/tvm/relax/testing/nn.py b/python/tvm/relax/testing/nn.py
index 184de0f862..a43dfab56e 100644
--- a/python/tvm/relax/testing/nn.py
+++ b/python/tvm/relax/testing/nn.py
@@ -34,6 +34,40 @@ def emit_te(func: Callable, *args: Any, **kwargs: Any) ->
relax.Var:
return relax.BlockBuilder.current().emit_te(func, *args, **kwargs)
+def _try_unique_name(name: str):
+ """Attempt to uniquify the name
+
+ If a `relax.BlockBuilder` is active, use it to return a unique
+ name. Otherwise, return the name itself.
+
+ Two distinct variables in Relax may have identical names.
+ However, for user readability, it is convenient to have all names
+ be unique within a Relax function. If a Placeholder or Parameter
+ is defined within an active `relax.BlockBuilder`, that context may
+ be used to provide a unique name. Otherwise, allow the duplicate
+ names.
+
+ Parameters
+ ----------
+ name: str
+
+ The variable name
+
+ Returns
+ -------
+ updated_name: str
+
+ The updated variable name
+
+
+ """
+ block_builder = relax.BlockBuilder.current()
+ if block_builder is None:
+ return name
+ else:
+ return block_builder.get_unique_name(name)
+
+
class Placeholder(relax.Var):
"""A placeholder variable that can represent model input."""
@@ -42,9 +76,7 @@ class Placeholder(relax.Var):
):
if not isinstance(shape, (list, tuple)):
raise TypeError("the shape of Placeholder is expected to be a list
or a tuple")
- super().__init__(
- relax.BlockBuilder.current().get_unique_name(name),
relax.TensorStructInfo(shape, dtype)
- )
+ super().__init__(_try_unique_name(name), relax.TensorStructInfo(shape,
dtype))
class Parameter(relax.Var):
@@ -55,9 +87,8 @@ class Parameter(relax.Var):
):
if not isinstance(shape, (list, tuple)):
raise TypeError("the shape of Parameter is expected to be a list
or a tuple")
- super().__init__(
- relax.BlockBuilder.current().get_unique_name(name),
relax.TensorStructInfo(shape, dtype)
- )
+
+ super().__init__(_try_unique_name(name), relax.TensorStructInfo(shape,
dtype))
class Module:
