Skip to content
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

feature: Jax ecosystem #1763

Closed
aarnphm opened this issue Jul 19, 2021 · 0 comments · Fixed by #3123
Closed

feature: Jax ecosystem #1763

aarnphm opened this issue Jul 19, 2021 · 0 comments · Fixed by #3123
Labels
help-wanted An issue currently lacks a contributor

Comments

@aarnphm
Copy link
Contributor

aarnphm commented Jul 19, 2021 •
edited
Loading

Is your feature request related to a problem? Please describe.

  • Supports some Jax-based frameworks, such as Flax, Haiku to BentoML

  • Add JaxArrayContainer

JAX support means BentoML should also support TPU and XLA for running inference.
-> we probably want to have a bentoml.io.JaxArray to fully utilize jax.lax.

Currently Jax implementation in BentoML (according to #3123) is a wrapper around jax.numpy interface.
@aarnphm aarnphm changed the title framework: Flax/Jax/Trax framework: Jax Jul 23, 2021
@aarnphm aarnphm closed this as completed Nov 10, 2021
@parano parano reopened this Apr 7, 2022
@ssheng ssheng added help-wanted An issue currently lacks a contributor and removed on-hold labels Jul 7, 2022
@aarnphm aarnphm changed the title framework: Jax framework: Jax, Flax, Trax Jul 9, 2022
This was linked to pull requests Oct 24, 2022
feat: cascade health checks to runner #3050
Merged
feat: Flax #3123
Merged
This was unlinked from pull requests Oct 24, 2022
feat: cascade health checks to runner #3050
Merged
feat: Flax #3123
Merged
@aarnphm aarnphm linked a pull request Oct 24, 2022 that will close this issue
feat: Flax #3123
Merged
@aarnphm aarnphm changed the title framework: Jax, Flax, Trax feature: Jax ecosystem Oct 24, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
help-wanted An issue currently lacks a contributor
Projects
None yet
Milestone
No milestone
Development

Successfully merging a pull request may close this issue.

feat: Flax aarnphm/BentoML
3 participants
@parano @ssheng @aarnphm