[FEATURE] JAX integration about zenml HOT 8 OPEN

AFAIK there is no builtin way to load/save JAX models, as it just contains the mathematical machinery for applying transformations and differentiating. For my own uses, I did export some models into HDF5, but that has little to do with JAX.

It's possible that first-party NN-libraries (flax, haiku, etc.) have some machinery for it, though. I think a small example on how to load/train/save a model might suffice, what do you have to implement for that? I'll see if I can come up with something if you point me to the necessary components :)

from zenml.

Totally. I got some time on the weekend, happy to take a look then. It may or may not be as straightforward for me though, depending on the usability of ZenML on M1. I'll get it done!

from zenml.

Haha @nicholasjng awesome to read this request. Love it. Will notify this thread as soon as we're ready for that undertaking! JAX is an awesome idea and thanks for the request. Lets do it 💪

from zenml.

Closing due to inactivity. We are migrating such issues to the roadmap for further voting :-)

from zenml.

Sorry to necro an old thread but I tried looking through the road map and could not find it. What's the status of this? Canceled, on hold, or completed?

Thank you :)

from zenml.

We didn't find enough demand so had to prioritize other things for now. However, I can resurface it if you're interested. Contributions are welcome here, what would a JAX integration look like? Similar to the tensorflow one, i.e., having the ability to pass (materialize) JAX models through ZenML pipelines?

from zenml.

@nicholasjng So sorry for the late reply , this slipped through. Happy to receive your contribution!

A good place to start would be to see the guide to add your own example. IMO, the example would resemble something like the LightGBM example.

You might need to implement a custom materializer to get this to work

I think this is a good starting point for JAX integration. WDYT?

from zenml.

M1 is a bit problem unfortunately. Let me know if it doesnt work though!

from zenml.