cc @U050CT4HR 👀

I am also interested in this!

But I don’t have a solution.

yeah I’m just exploring/scoping out/brainstorming atm 🙂

I think its a good idea and I would be interested in trying something like this What type of wrapper do you want to make? e.g. • a partial wrapper (ie make it easier to use jax from Clojure) but you still need to understand JAX • A full wrapper (make a new Clojure library on top of JAX) - user doesn’t need to understand JAX. I think there are a couple of things to bare in mind if you attempt a wrapper: • JAX is such a large library I imagine that any attempt to wrap it will end up as a leak abstraction layer. That’s not necessarily bad, but worth attempting to address in the design. • Clojure JAX users (via lib-pythonclj) can currently follow along with Python tutorials - will that still be possible with the wrapper or will you need to create new tutorial/training material.

I actually haven’t started working w/JAX from lib-pythonclj yet (just Python so far). I’d be curious what that experience is like? I have it on a list of things to look into. But curious if others that have done so have specific pain points that are important to them. I.e. I’d want to make sure a clojure wrapper or jaxpr compiler did something better and solved problems. It’s also possible that a JAX inspired framework would be a better fit, but this would take more resources and be a riskier investment of effort. (but high risk/reward, likely).

I've used some other ML libraries with libpython-clj (e.g. transformers) and you can basically write Python code in CLojure (e.g. one line of Python = one line of Clojure). That lets you get to parity - then you can rewrite the code to make it more idiomatic.

I tend as well to think that "wrapping a python library" is not a good idea. The wrapped library will change faster then you can adapt the wrapper. I tend as well to keep the "Interop code" I need for my concrete problem in a single ns, and give it a nicer Clojure intercae, using pure Clojure data.

Maybe an other word of warning. I have seen that libpython-clj works well with 99% of python libraries. The 1% it fails upon are the ones which try hard to be "fast"... Machine learning libraries might fall in this area. So any "wrapping project" of a ML library has a certain risk that it is in the 1 % on which libpython-clj will fail. Libpython-clj uses the "embed python interpreter" concept with is not 100% compatible with CPython (the native python interpreter which nearly everybody is using) Reporting bugs in this will be answered by "use CPython"...

yeah, I’m also curious to what degree there’s a path forward in “inspired by jax,” or emitting jaxpr directly. it just hurts me to go from a language w/immutability by default, pure functions as good style, and a jit to a much more simply interpreted language with a GIL so I can get a lift from a framework in said language that leverages immutability via jit, applied through decorators :melting_face: whether wrapping or switching tools as a human. I would like to eliminate the middle man here.