Skip to content

Jax custom call #12396

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
Closed
wants to merge 19 commits into from
Closed

Jax custom call #12396

wants to merge 19 commits into from

Conversation

@zhangqiaorjc
Copy link
Collaborator

@zhangqiaorjc zhangqiaorjc commented Sep 16, 2022

No description provided.

@zhangqiaorjc zhangqiaorjc marked this pull request as draft September 16, 2022 21:37
hawkinsp
hawkinsp reviewed Sep 16, 2022
View reviewed changes
jaxlib/jax_custom_call.h Outdated
JaxCustomCallStatus* status;
};

void JaxCustomCallStatusSetSuccess(JaxCustomCallStatus* status);
Copy link
Collaborator

@hawkinsp hawkinsp Sep 16, 2022

Choose a reason for hiding this comment

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

BTW, linkage may be hard here. My suggestion to Sharad was that we pass function pointers for all API entry points a custom call might need (i.e., like a c++ class implemented in C).

Copy link
Collaborator Author

@zhangqiaorjc zhangqiaorjc Sep 21, 2022

Choose a reason for hiding this comment

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

Updated to use runtime linkage instead of link-time linkage a la Numpy import_array() style.

@sharadmv sharadmv force-pushed the jax_custom_call branch 3 times, most recently from a78db4a to a379de5 Compare September 20, 2022 00:33
@sharadmv sharadmv mentioned this pull request Oct 4, 2022
Closed
@zhangqiaorjc zhangqiaorjc marked this pull request as ready for review October 26, 2022 21:59
@sharadmv sharadmv added the pull ready Ready for copybara import and testing label Oct 26, 2022
@zhangqiaorjc zhangqiaorjc force-pushed the jax_custom_call branch 2 times, most recently from a18f38c to 17fb10c Compare October 26, 2022 22:36
@zhangqiaorjc
Copy link
Collaborator Author

zhangqiaorjc commented Oct 26, 2022

We are importing this to explore some things. @hawkinsp hold off on reviews. Thanks!

@jlu-spins
Copy link

jlu-spins commented Jan 23, 2023

Hi! This looks really useful. Is anyone still committed to making this happen?

@zhangqiaorjc
Copy link
Collaborator Author

zhangqiaorjc commented Jan 23, 2023

@jlu-spins yes, we are making an improved version with better C++ typed argument support. But it creates a new dependency on some XLA component that's still WIP. We are actively working on this, and will hopefully have a version in a few weeks.

In the meanwhile, https://github.com/dfm/extending-jax is still a good guide for doing JAX custom call at head.

@jlu-spins
Copy link

jlu-spins commented Jan 23, 2023

@zhangqiaorjc Thank you!

@MoFHeka
Copy link

MoFHeka commented Jul 23, 2024

How create a custom class instance with several call method and its own stateful resource. Something like Extending TorchScript with Custom C++ Classes

@dfm
Copy link
Contributor

dfm commented Jul 24, 2024

I'm going to close this PR since this work has been superseded by XLA's FFI API which is documented in #22095.

@MoFHeka — Please feel free to open a new discussion topic with more details about your specific use case if it would be useful to continue the conversation.

@dfm dfm closed this Jul 24, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@hawkinsp hawkinsp hawkinsp left review comments

Assignees

@zhangqiaorjc zhangqiaorjc

@sharadmv sharadmv

@cky9301 cky9301

Labels

pull ready Ready for copybara import and testing

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

7 participants

@zhangqiaorjc @jlu-spins @MoFHeka @dfm @hawkinsp @sharadmv @cky9301