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
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
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
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
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
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.
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
Wanted to share this library that I'm open-sourcing now after working on it for the past year.
You write code like in JAX/NumPy, but it’s fully interactive on the frontend and compiles down to shaders on the user’s GPU (with WebGPU). So far I’ve used it for purely frontend-only ML demos. https://jax-js.com/mobileclip
Thanks! I think that's a great idea. Containers are nice, but there's nothing quite like a "real" VM system in terms of learning + developing, and they can be very cheap if you're OK with the performance hit from nested virtualization :)
Another idea I had is to customize the rootfs. It boots Alpine right now, but should be trivial to change it to anything else (and it only takes ~5 seconds to build!).
Hello! Thrilled to be sharing something I've worked on for the past 8 months at Modal — a cloud, GPU notebooks product.
We wanted to improve the Jupyter Notebook experience while offering high-performance, cloud compute. So we sought out to fix some of the problems in cloud offerings for notebooks, and we ended up rebuilding a lot of the interface ourselves. This means you can do things like:
- Within the same notebook, swap from CPU to GPU, and the kernel will restart in seconds while keeping your work.
- Use _lots_ of hardware: up to 8x H100/B200 GPUs, not just one.
- Automatic idle shutdown (configurable timeout) so you don't get a huge bill / forget to stop your notebook.
- Real-time collaborative editing.
- Integrations with Pyright and Ruff for a better IDE-like experience.
Unlike Google Colab, we are a full cloud offering, so you can integrate with the rest of the Modal platform — use custom images, attach persistent volumes and so on. We've heard these are the actual, boring but really important features you need to do experimental work in notebooks.
Happy to hear any feedback! Would love for you to try it out.