LLM for writing JAX-compatible code

[cool-RR]

Hi everyone,

I've been working with JAX for 6 months now and I'm liking it. One of the challenges in using JAX is learning to write logic in JAX-compatible ways, like using jnp.where instead of if. Sometimes I'm successful at doing that, and sometimes not. I often ask ChatGPT and Claude to convert code for me from old-fashioned Python to JAX-compatible. Sometimes they get it right and sometimes I have to do some trial-and-error with the results that they give me until I get it right. I guess the major LLMs have some knowledge of JAX but not a lot.

It would be cool if there was an LLM that had better knowledge of JAX so it could do this converting work better. I wonder whether RAG is the answer, or maybe there should be a JAX cheatsheet I could upload to ChatGPT/Claude/Gemini along with my request so it'll know how to answer JAX questions.

Maybe someone has already thought of something like this and built it? Maybe someone would be interested in building it?

[Qazalbash]

A naive question, have you tried using JAX documentation for training an LLM or could this idea work?

[cool-RR]

No, but I have tried to do this in the past for a different open-source project and I didn't get good results.