Show HN: Jax-JS, array library in JavaScript targeting WebGPU
1–10 of 23 posts
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#2What is the state of web ML? Anybody doing cool things already? How about https://www.w3.org/TR/webnn/ ?
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#3I have a project using tfjs and jax-js is very exciting alternative. However during porting I struggle a lot with `.ref` and `.dispose()` API. Coming from tfjs where you garbage collect with `tf.tidy(() => { ... })`, API in jax-js seems very low-level and error-prone. Is that something that can be improved or is it inherent to how jax-js works?
Would `using`[0] help here?
[0]: https://developer.mozilla.org/en-US/docs/Web/JavaScript/Refe...
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#4This project is an inspiration, I've been working on porting tinygrad to [Lean](github.com/alok/tinygrad)
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#5Hey Eric, great to see you've now published this! I know we chatted about this briefly last year, but it would be awesome to see how the performance of jax-js compares against that of other autodiff tools on a broader and more standard set of benchmarks: https://github.com/gradbench/gradbench
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#6This is really great. I don't do ML stuff. But I some mathy things that would benefit from running in the GPU so it's great to see the Web getting this.
I hope this will help grow the js science community.
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#7Hey Eric, great to see you've now published this! I know we chatted about this briefly last year, but it would be awesome to see how the performance of jax-js compares against that of other autodiff tools on a broader and more standard set of benchmarks: https://github.com/gradbench/gradbench
For sure! It looks like this is benchmarking the autodiff cpu time, not the actual kernels though, which (correct me if I’m wrong) isn’t really relevant for an ML library — it’s more for if you have a really complex scientific expression
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#8What is the state of web ML? Anybody doing cool things already? How about https://www.w3.org/TR/webnn/ ?
onnx on the web has the most models available and can use webgpu which is available everywhere.
Huggingface’s transformers.js uses it. And I use that for https://workglow.dev (also tensorflow mediapipe though that is using wasm).
I don’t think webnn has gone anywhere and is too restrictive.
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#9I have a project using tfjs and jax-js is very exciting alternative. However during porting I struggle a lot with `.ref` and `.dispose()` API. Coming from tfjs where you garbage collect with `tf.tidy(() => { ... })`, API in jax-js seems very low-level and error-prone. Is that something that can be improved or is it inherent to how jax-js works? Would `using`[0] help here? [0]: https://developer.mozilla.org/en-US/docs…
I don’t think tf.tidy() is a sound API under jvp/grad transformations, also it prevents you from using async which makes it incompatible with GPU backends (or blocks the page), a pretty big issue. https://github.com/tensorflow/tfjs/issues/5468
Thanks for the feedback though, just explaining how we arrived at this API. I hope you’d at least try it out — hopefully you will see when developing that the refs are more flexible than alternatives.
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#10What is the state of web ML? Anybody doing cool things already? How about https://www.w3.org/TR/webnn/ ?
onnx on the web has the most models available and can use webgpu which is available everywhere. Huggingface’s transformers.js uses it. And I use that for https://workglow.dev (also tensorflow mediapipe though that is using wasm). I don’t think webnn has gone anywhere and is too restrictive.
Since ONNX is just a model data format, you can actually parse and run ONNX files in jax-js as well. Here’s an example of running DETR ResNet-50 from Xenova’s transformers.js checkpoint in jax-js
https://jax-js.com/detr-resnet-50
I don’t think I intend to support everything in ONNX right now, especially quant/dequant, but eventually it would be interesting to see if we can help accelerate transformers.js with a jax-js backend + goodies like kernel fusion
jax-js is more trying to explore being an ML research library, rather than ONNX which is a runtime for exported models