slyubomirsky commented on code in PR #15577:
URL: https://github.com/apache/tvm/pull/15577#discussion_r1307943347
##########
python/tvm/relax/struct_info.py:
##########
@@ -42,14 +43,68 @@ class PrimStructInfo(StructInfo):
Parameters
----------
- dtype : str
- The data type of the prim value.
+ dtype_or_expr : Union[str, DataType, PrimExpr]
+
+ The data type of the prim value, or a known expression for the prim
+ value.
"""
+ value: Optional[PrimExpr]
dtype: str
- def __init__(self, dtype: str, span: Span = None) -> None:
- self.__init_handle_by_constructor__(_ffi_api.PrimStructInfo, dtype,
span) # type: ignore
+ def __init__(
+ self,
+ dtype: Optional[Union[str, DataType]] = None,
+ value: Optional[Union[int, float, PrimExpr]] = None,
+ span: Span = None,
+ ) -> None:
+ # Guard against incorrect usage. For backwards compatibility,
+ # the dtype and value are in the opposite order from most
+ # usages. While PrimStructInfo could take a single positional
+ # argument and check the type, this would require an API
+ # difference from TVMScript's PrimProxy, which cannot.
+ # (PrimProxy uses string arguments for datatype, and also for
+ # inline variable definitions when used in a function
+ # signature, and requires separate arguments to distinguish
+ # the two cases.)
Review Comment:
Very user-friendly touch :)
##########
tests/python/relax/test_bind_symbolic_vars.py:
##########
@@ -183,6 +183,91 @@ def expected(A: R.Tensor(["M", 16])):
tvm.ir.assert_structural_equal(expected, after)
+def test_bind_symbolic_vars_in_shape_expr():
+ """The bound variable should be replaced when appearing in R.Shape"""
+
+ @R.function(private=True)
+ def before(A: R.Tensor(["M * N"]), x: R.Shape(["M", "N"])):
Review Comment:
Are we sure we want to permit cases like this? I believe I wrote the
specification draft to require argument vars to be in binding positions, left
to right. If we permit uses before the binding position, then should the
type-checker scan through the arguments, identify the bindings first, and only
then check the position?
--
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]
