docs: initial prototype of exporting Lux models to Jax

[avik-pal]

fixes #453

[avik-pal]

Tested locally that this works (needs a custom build of EnzymeJAX)

[avik-pal]

cc @wsmoses if you have any suggestions

[wsmoses]

Incidentally we can probably make a flag to compile/jit/code hlo to not transpose the inputs

[wsmoses]

But maybe regardless it’s worth explaining why: Julia default uses col major vs jax default uses row major

[avik-pal]

Incidentally we can probably make a flag to compile/jit/code hlo to not transpose the inputs

I feel this will be very confusing for people who use Python exclusively. For people who use both (Julia and Python), switching is not that hard when converting the model.

I think we can add a function that serializes all the inputs from Julia + mlir code and we have a python function to deserialize it. That would also be handy for pre-trained models and such

But maybe regardless it’s worth explaining why: Julia default uses col major vs jax default uses row major

Agreed

[wsmoses]

Extra fun fact, hypothetically if the model weights themselves weren’t traced, it should output MLIR containing the weights and thus one can export inference of a pre trained model too!

[avik-pal]

This is actually a neat way to do it, but will need some additional plumbing on Lux end. Currently if luxlib sees mismatch in devices (CPUDevice vs ReactantDevice) it throws an error. This is very handy if users forgot to move something to GPU and would previously hit a cryptic error deep in the stack.

With Reactant it might be worth reworking some of those things

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 37.18%. Comparing base (3986545) to head (b52f96c).
Report is 2 commits behind head on main.

❗ There is a different number of reports uploaded between BASE (3986545) and HEAD (b52f96c). Click for more details.

HEAD has 56 uploads less than BASE

Flag	BASE (3986545)	HEAD (b52f96c)
	57	1

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #1088       +/-   ##
===========================================
- Coverage   83.79%   37.18%   -46.62%     
===========================================
  Files         147       61       -86     
  Lines        6036     2883     -3153     
===========================================
- Hits         5058     1072     -3986     
- Misses        978     1811      +833     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.