Skip to content

Add input output aliases to jax FFI. #787

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 9 commits into
base: main
Choose a base branch
from
Open

Add input output aliases to jax FFI. #787

wants to merge 9 commits into from

Conversation

btaba
Copy link
Contributor

@btaba btaba commented Jun 14, 2025 •
edited
Loading

Description

Adds input output aliases (ref) to the JAX interop code.

Before your PR is "Ready for review"

  • All commits are signed-off to indicate that your contribution adheres to the Developer Certificate of Origin requirements
  • Necessary tests have been added
  • Documentation is up-to-date
  • Auto-generated files modified by compiling Warp and building the documentation have been updated (e.g. __init__.pyi, functions.rst)
  • Code passes formatting and linting checks with pre-commit run -a

btaba added 2 commits June 13, 2025 17:07
@btaba
Add input output aliases to jax FFI.
fdd066f 
Signed-off-by: Baruch Tabanpour <[email protected]>
@btaba
ruff.
0903746 
Signed-off-by: Baruch Tabanpour <[email protected]>
@shi-eric shi-eric requested a review from nvlukasz June 14, 2025 00:10
jaro-sevcik
jaro-sevcik reviewed Jun 21, 2025
View reviewed changes
Copy link
Contributor

@jaro-sevcik jaro-sevcik left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM modulo some non-blocking nits.

@btaba
Copy link
Contributor Author

btaba commented Jun 23, 2025

Thanks for the review @jaro-sevcik !

btaba and others added 5 commits June 23, 2025 11:13
@btaba
Change semantics of if statement.
f681569 
Signed-off-by: Baruch Tabanpour <[email protected]>
@shi-eric @btaba
Enhance error reporting for @wp.func calls in Python scope (NVIDIAGH-521
dcbb310 
)

Improve contextual error messages for function overload failures.

Update argument type inference documentation.

Signed-off-by: Eric Shi <[email protected]>
Signed-off-by: Baruch Tabanpour <[email protected]>
@shi-eric @btaba
Add Lee et al.
c51d333 
Signed-off-by: Eric Shi <[email protected]>
Signed-off-by: Baruch Tabanpour <[email protected]>
@btaba
Remove in_out_argnames from FfiKernel.
e71e3c5 
Signed-off-by: Baruch Tabanpour <[email protected]>
@btaba
In out argnames
6bce4a2 
Signed-off-by: Baruch Tabanpour <[email protected]>
@btaba
ruff
b7ec880 
Signed-off-by: Baruch Tabanpour <[email protected]>
@btaba
Copy link
Contributor Author

btaba commented Jun 24, 2025

Any remaining blockers?

@shi-eric
Copy link
Contributor

Any remaining blockers?

GitHub is showing that there are merge conflicts in your branch

@btaba
Copy link
Contributor Author

btaba commented Jun 25, 2025

Thanks @shi-eric, sorry I missed that. Just merged main into this branch.

@shi-eric shi-eric added this to the 1.9.0 milestone Jun 25, 2025
@shi-eric
Copy link
Contributor

Thanks @shi-eric, sorry I missed that. Just merged main into this branch.

Thanks! @nvlukasz will take a final pass over this request this week before we merge it.

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

@jaro-sevcik jaro-sevcik jaro-sevcik left review comments

@nvlukasz nvlukasz Awaiting requested review from nvlukasz

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
1.9.0
Development

Successfully merging this pull request may close these issues.
3 participants
@btaba @shi-eric @jaro-sevcik