Jax code generation (for instance openCL)

[EelcoHoogendoorn]

One thing id like to do is write code in JAX and compile these JAX functions to a minimal-dependency openCL program, with the intent to run them on embedded systems. (Kinda like tinygrad does, but with composible transformations, and without the horrible pytorch syntax.)

It seems to me that doing such codegen from a JAX expression tree should be very possible... but im not in the mood to reinvent the wheel and my google-fu is coming up entirely empty handed when searching for any such thing (wether its openCL or CUDA or vulcan or plain C).

Am I just not searching right? Is there some fundamental reason im missing why this actually wouldnt be possible to add such alternative codegen options, because of some deep integration with XLA? Or would it actually not be all that hard, but it just hasnt been done before?