Show HN: Jax-JS, array library in JavaScript targeting WebGPU
11–20 of 23 posts
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#12Could not run the demos on Firefox. On Chromium, the Great Expectations loads but then nothing happens.
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#13Hey 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
#14Earlier quoted context omitted.
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
Nope, both are measured! In fact, the time to do the autodiff transformation isn't even reflected in the charts shown on the README and the website; those charts only show the time to actually run the computations.
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#15Earlier quoted context omitted.
Nope, both are measured! In fact, the time to do the autodiff transformation isn't even reflected in the charts shown on the README and the website; those charts only show the time to actually run the computations.
Hm okay, seems like an interesting set of benchmarks — let me know if there’s anything I can do to help make jax-js more compatible with your docker setup
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#16I will be following this project closely! Best of luck Eric! Do you have plans to keep working on it for sometime? Is it a side project or will you abe ble to commit to jax-js longer term?
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#17Congrats on the launch! This is a very exciting project because the only decent autodiff implementation in typescript was tensorflowjs, which has been completely abandonned by Google. Everyone uses onnx runtime web for inference but actually computing gradients in typescript was surprisingly absent from the ecosystem since tfjs died. I will be following this project closely! Best of luck Eric! Do you have plans to ke…
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#18Earlier quoted context omitted.
Hm okay, seems like an interesting set of benchmarks — let me know if there’s anything I can do to help make jax-js more compatible with your docker setup
It should be fairly straightforward; feel free to open a PR following the instructions in CONTRIBUTING.md :)
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#19Earlier quoted context omitted.
It should be fairly straightforward; feel free to open a PR following the instructions in CONTRIBUTING.md :)
I don’t think this is straightforward but it may be a skill issue on my part. It would require dockerizing headless Chrome with WebGPU support and dynamically injecting custom bundled JavaScript into the page, then extracting the results with Chrome IPC
Re: Show HN: Jax-JS, array library in JavaScript targeting WebGPU
#20I 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 fle…