Show HN: Jax-JS, array library in JavaScript targeting WebGPU

https://ss.ekzhang.com/p/jax-js-an-ml-library-for-the-web

84•ekzhang•1mo ago

Comments

esafak•1mo ago

What is the state of web ML? Anybody doing cool things already? How about https://www.w3.org/TR/webnn/ ?

sroussey•1mo ago

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•1mo ago

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•1mo ago

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•1mo ago

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.

mlajtos•1mo ago

I'll grind jax-js more and see if refs become invisible then. Thanks for a great project!

yuppiemephisto•1mo ago

This project is an inspiration, I've been working on porting tinygrad to [Lean](github.com/alok/tinygrad)

sestep•1mo ago

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•1mo ago

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•1mo ago

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•1mo ago

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•1mo ago

It should be fairly straightforward; feel free to open a PR following the instructions in CONTRIBUTING.md :)

ekzhang•1mo ago

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

sestep•1mo ago

Ahh no you're right, I forgot about the difficulties for GPU specifically; apologies for my overly curt earlier message. More accurately: I think this is definitely possible (Troels and I have talked a bit about this previously) and I'd be happy to work together if this is something you're interested in. I probably won't work on this if you're not interested on your end, though.

bobajeff•1mo ago

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•1mo ago

Could not run the demos on Firefox. On Chromium, the Great Expectations loads but then nothing happens.

ekzhang•1mo ago

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

forgotpwd16•1mo ago

According to page WebGPU supported (`dom.webgpu.enabled` flag) but is only enabled by default on Windows & macOS (i.e. not Linux).

fouronnes3•1mo ago

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•1mo ago

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

forgotpwd16•1mo ago

Very nice work. Like how it supports webgpu but also cpu/wasm/webgl. Would love to read more on the internals & design choices made like e.g. ref counting in README.

P.S. And thanks for taking your time working on this and releasing something polished rather a Claude slop made within few days as seems to be the norm now.

sbondaryev•1mo ago

The examples are great. It would be really nice to have a sandbox with the full training code (e.g. MNIST) to play with.

Show HN: LocalGPT – A local-first AI assistant in Rust with persistent memory

Haskell for all: Beyond agentic coding

SectorC: A C Compiler in 512 bytes (2023)

Speed up responses with fast mode

Software factories and the agentic moment

Brookhaven Lab's RHIC concludes 25-year run with final collisions

Vouch

Do you have a mathematically attractive face?

Stories from 25 Years of Software Development

Hoot: Scheme on WebAssembly

IBM Beam Spring: The Ultimate Retro Keyboard

First Proof

FDA intends to take action against non-FDA-approved GLP-1 drugs

Vocal Guide – belt sing without killing yourself

Show HN: I saw this cool navigation reveal, so I made a simple HTML+CSS version

Al Lowe on model trains, funny deaths and working with Disney

Show HN: Axiomeer – An open marketplace for AI agents

LLMs as the new high level language

Show HN: A luma dependent chroma compression algorithm (image compression)

Start all of your commands with a comma (2009)

The AI boom is causing shortages everywhere else

Microsoft account bugs locked me out of Notepad – Are thin clients ruining PCs?

The F Word

I write games in C (yes, C) (2016)

Selection rather than prediction

Eigen: Building a Workspace

The silent death of good code

Reinforcement Learning from Human Feedback

OpenCiv3: Open-source, cross-platform reimagining of Civilization III

Learning from context is harder than we thought

Show HN: LocalGPT – A local-first AI assistant in Rust with persistent memory

Haskell for all: Beyond agentic coding

SectorC: A C Compiler in 512 bytes (2023)

Speed up responses with fast mode

Software factories and the agentic moment

Brookhaven Lab's RHIC concludes 25-year run with final collisions

Vouch

Do you have a mathematically attractive face?

Stories from 25 Years of Software Development

Hoot: Scheme on WebAssembly

IBM Beam Spring: The Ultimate Retro Keyboard

First Proof

FDA intends to take action against non-FDA-approved GLP-1 drugs

Vocal Guide – belt sing without killing yourself

Show HN: I saw this cool navigation reveal, so I made a simple HTML+CSS version

Al Lowe on model trains, funny deaths and working with Disney

Show HN: Axiomeer – An open marketplace for AI agents

LLMs as the new high level language

Show HN: A luma dependent chroma compression algorithm (image compression)

Start all of your commands with a comma (2009)

The AI boom is causing shortages everywhere else

Microsoft account bugs locked me out of Notepad – Are thin clients ruining PCs?

The F Word

I write games in C (yes, C) (2016)

Selection rather than prediction

Eigen: Building a Workspace

The silent death of good code

Reinforcement Learning from Human Feedback

OpenCiv3: Open-source, cross-platform reimagining of Civilization III

Learning from context is harder than we thought

Show HN: Jax-JS, array library in JavaScript targeting WebGPU

Comments