TE
TechEcho
Home24h TopNewestBestAskShowJobs
GitHubTwitter
Home

TechEcho

A tech news platform built with Next.js, providing global tech news and discussions.

GitHubTwitter

Home

HomeNewestBestAskShowJobs

Resources

HackerNews APIOriginal HackerNewsNext.js

© 2025 TechEcho. All rights reserved.

How the jax.jit() JIT compiler works in jax-js

4 pointsby ekzhang10 days ago

1 comment

ekzhang10 days ago
Hello - I wanted to share a bit about an open-source project I’ve been working on, building a fundamental library for ML and numerical computing in the browser.<p>There are different tradeoffs here, but I’d like to replicate the speed and ergonomics of Python libraries with NumPy&#x2F;PyTorch&#x2F;JAX without the difficulty of distribution. jax-js is lightweight enough to be embedded in a website, but optimized to take advantage of tech like Wasm and WebGPU.<p>Let me know if you have any feedback on this compiler design. It’s working well and faster than tfjs at matrix multiplication, while having kernel fusion as well.