Show HN: Jax-JS, array library in JavaScript targeting WebGPU
4 months ago
- #javascript
- #web-development
- #machine-learning
- jax-js是一个面向网页的机器学习库，用纯JavaScript重新实现了Google DeepMind的JAX框架
- 它通过WebGPU和WebAssembly在浏览器中运行，提供接近原生的性能，实现了前端数值计算
- 该库提供类似JAX的API，包括自动微分(grad)、向量化(vmap)和即时编译(jit)等功能
- jax-js支持在浏览器中训练神经网络，例如MNIST数据集可在数秒内达到>99%的准确率
- 它利用WebGPU进行GPU加速，为矩阵乘法等操作生成优化内核
- 该项目是开源的，可通过npm获取，零依赖且完全采用JavaScript实现
- 性能表现优异，在现代硬件上矩阵乘法运算可达>3 TFLOP/s
- 该库包含热模块重载等功能，支持模型训练期间实时编辑代码
- 未来改进可能包括WebAssembly优化、WebGL回退方案以及更多JAX功能支持
- 该项目展示了完整机器学习框架在浏览器中运行的潜力，超越了ONNX等传统模型运行时