Add jax compatible api

[sky-2002]

This PR adds a JAX compatible API, refer issue #1027

Looking forward for review and open to feedback(especially in writing tests).

[lapp0]

Great work so far!

Could you please add jax to benchmarks/bench_processors.py as well?

Also could you fill me in on how you intend to use a jax based model, just so I have context on how Outlines is being used here?

Thanks!

[sky-2002]

Hey @lapp0 , thanks for the feedback.
Have added to benchmarks(though I haven't run those yet, will possibly do)

Regarding intent to use with JAX, @borisdayma(author of issue mentioned above) can help us understand jax based model usage with outlines.

PS: I currently have no use case with JAX, but I was trying out outlines and found it very interesting, so wanted to contribute.

[lapp0]

PS: I currently have no use case with JAX, but I was trying out outlines and found it very interesting, so wanted to contribute.

Thanks for clarifying and your interest in the project.

Have added to benchmarks(though I haven't run those yet, will possibly do)

I'll add the run-benchmarks label, which will trigger their run once the workflow is approved.

Regarding intent to use with JAX, @borisdayma(author of issue mentioned above) can help us understand jax based model usage with outlines.

@borisdayma could you review and smoke test to ensure this fits your desired use case?

[lapp0]

no need for this if you use the below comment

[borisdayma]

This looks great but I’ll still need some time to try it out.

The goal is to use it with repo’s such as:

• https://github.com/AI-Hypercomputer/maxtext with LLM’s in JAX (llama, gemma, etc)
• https://github.com/borisdayma/clip-jax for VLM sampling

Since the basic tests seem to work I expect it should be fine and could create a new issue if for some reason it prevents JAX compilation.

[rlouf]

Thank you for contributing and addressing @lapp0's comments. It looks like we'll be good to merge if the CI turns green.