[HN Gopher] Show HN: Jax-JS, array library in JavaScript targeti...
___________________________________________________________________
Show HN: Jax-JS, array library in JavaScript targeting WebGPU
Author : ekzhang
Score : 73 points
Date : 2026-01-06 18:19 UTC (4 hours ago)
(HTM) web link (ss.ekzhang.com)
(TXT) w3m dump (ss.ekzhang.com)
| esafak wrote:
| What is the state of web ML? Anybody doing cool things already?
| How about https://www.w3.org/TR/webnn/ ?
| sroussey wrote:
| 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.
| ekzhang wrote:
| 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
| mlajtos wrote:
| I 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...
| ekzhang wrote:
| 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.
| yuppiemephisto wrote:
| This project is an inspiration, I've been working on porting
| tinygrad to [Lean](github.com/alok/tinygrad)
| sestep wrote:
| Hey 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
| ekzhang wrote:
| 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
| sestep wrote:
| 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.
| ekzhang wrote:
| 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
| sestep wrote:
| It should be fairly straightforward; feel free to open a
| PR following the instructions in CONTRIBUTING.md :)
| bobajeff wrote:
| This 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.
| maelito wrote:
| Could not run the demos on Firefox. On Chromium, the Great
| Expectations loads but then nothing happens.
| ekzhang wrote:
| Firefox doesn't support WebGPU yet, you can run programs in the
| REPL through other backends like Wasm/WebGL: https://jax-
| js.com/repl
|
| See: https://caniuse.com/webgpu
| fouronnes3 wrote:
| Congrats 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 keep working on it for sometime? Is it a side
| project or will you abe ble to commit to jax-js longer term?
| ekzhang wrote:
| Yes, we are actively working on it! The goal is to be a full ML
| research library, not just a model inference runtime. You can
| join the Discord to follow along
___________________________________________________________________
(page generated 2026-01-06 23:09 UTC)