Supporting JAX-inspired WebML frameworks/libraries

[josephrocca]

I'm not qualified to write an actual proposal here so this is just a placeholder issue for discussion about supporting JAX-inspired JS frameworks. I originally created an issue in the WebNN repo was advised by @anssiko to create an issue here instead.

The thrust of the original issue was that JAX is becoming more popular, and that as the foundations of WebNN/WebML are built, it may be important to take into account its growing popularity so that highly-performant "JAX.js" type frameworks are possible in the future.

In the original issue I said:

IIRC, WebNN's initial focus is on inference for networks trained using non-web frameworks, which makes sense, but this question is more about the long-term trends, given the seeming possibility of JAX-like frameworks becoming the norm.

But I'd also like to add that gradients are required for the "guiding" done by models like VQGAN+CLIP. It seems like guiding embeddings/latents/inputs via models like CLIP is becoming more popular. It may end up being important for WebML to support this type of "inference", rather than just purely forward-prop inference, but maybe it's too early to say. Either way, allowing for the possibility of this when designing the foundations seems like a good idea.

Links:

• Initial discussion of original issue: https://www.w3.org/2021/10/14-webmachinelearning-minutes.html#t02
• What is JAX? https://github.com/google/jax#what-is-jax

[anssiko]

Thanks @josephrocca. Adding @wchao1115 for comments.

FYI @nsthorat and @dsmilkov.