fix(CuTeDSL): preserve wrapped signatures for type checkers - #3528
Open
iwzbi wants to merge 1 commit into
Open
Conversation
dsl_user_op is applied to every user-facing CuTe DSL wrapper (local_tile,
make_tiled_mma, gemm, etc.) but its annotation erased the wrapped callable
to Callable[..., Any]. Static analyzers (pyright/basedpyright/mypy) and
language servers therefore see only (*args, **kwargs): editors show no
parameter names in signature help and IDEs report bogus unknown-keyword
calls as valid.
Use a TypeVar bound to Callable so the decorator returns the exact
signature of the wrapped function, following the recommended decorator
pattern in the pyright documentation:
_FuncT = TypeVar('_FuncT', bound=Callable[..., Any])
def decorator(func: _FuncT) -> _FuncT: ...
Verified with pyright: cute.local_tile(mA, tiler=..., coord=...) now
resolves its real named parameters, while cute.local_tile(mA, xxxxx=...)
is correctly rejected as 'No parameter named "xxxxx"'.
Runtime behavior is unchanged (functools.wraps was already applied).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
@dsl_user_opdecorates every user-facing CuTe DSL wrapper (local_tile,make_tiled_mma,copy,gemm, etc.), but its annotation wasCallable[..., Any] -> Callable[..., Any], which erases the wrapped callable's signature for static analyzers.Pyright/basedpyright (and language servers built on them) therefore see only
(*args, **kwargs). In practice this means:cute.local_tile(yields notiler/coordhints;cute.local_tile(mA, xxxxx=...)passes type checking);@dsl_user_opwrappers (vLLM, SGLang, FlashInfer) are affected the same way.Change
Make the decorator signature-preserving with a
TypeVarbound toCallable, the pattern recommended by pyright's documentation:Runtime behavior is unchanged (
functools.wrapswas already applied to the wrapper).# type: ignore[return-value]documents that the innerwrapperintentionally has a generic signature but must satisfy_FuncT.Verification
python -m py_compile op.py— syntax OK.cute.local_tile(mA, tiler=(128, 64), coord=(0, 0))type-checks;cute.local_tile(mA, xxxxx=(128, 64))is rejected withNo parameter named "xxxxx", andArguments missing for parameters "tiler", "coord"is reported for incomplete calls.Fixes editor/signature-help blindness for the entire public CuTe DSL surface.