Skip to content

Add missing defaults to tri, tril, triu, gather_mm signatures - #4492

Open
rishabhsai wants to merge 1 commit into
ml-explore:mainfrom
rishabhsai:fix-python-sig-defaults
Open

rishabhsai wants to merge 1 commit into
ml-explore:mainfrom
rishabhsai:fix-python-sig-defaults

Conversation

@rishabhsai

Copy link
Copy Markdown

The signatures for tri, tril, triu, and gather_mm omit defaults that their Python bindings already provide. The generated stubs therefore mark these arguments as required, causing Pyright to reject valid calls such as mx.tri(3) and mx.gather_mm(x, x).

This updates the four nb::sig strings in python/src/ops.cpp to match the existing bindings and docstrings:

  • tri: declare m: int | None = None and k: int = 0.
  • tril and triu: declare k: int = 0.
  • gather_mm: declare both index arguments as array | None = None.

There is no runtime behavior change.

Validation used a CPU-only source build. The regenerated stubs contain the corrected signatures, and Pyright reports zero errors on six example calls that previously produced six missing-argument errors. Those calls execute successfully both before and after the change. The existing tri, tril, and triu tests pass, as does the clang-format check.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants