jax-js: 웹용 머신러닝 라이브러리 및 컴파일러

jax-js는 JavaScript로 작성된 브라우저용 머신러닝 라이브러리이자 컴파일러다. WebGPU와 WebAssembly 커널을 JIT 컴파일해 신경망, 이미지 알고리즘, 시뮬레이션, 수치 연산을 클라이언트 환경에서 실행한다.

Chrome, Firefox, iOS, Android 등 주요 브라우저에서 동작하며, 외부 의존성 없이 npm install @jax-js/jax로 프로젝트에 추가할 수 있다. NumPy·JAX와 유사한 API를 제공하고 JAX 및 PyTorch에서 영감을 받은 설계를 적용했다.

주요 특징은 다음과 같다.

• 브라우저에서 완전히 클라이언트 사이드로 실행
• WebGPU, WebGL, WebAssembly 지원
• 고수준 연산을 WebGPU·WebAssembly 커널로 변환하는 경량 컴파일러
• GPU와 CPU를 활용한 로컬 LLM 추론
• AI 모델, 시뮬레이션, 기하 처리, 수치해석 관련 인터랙티브 예제 제공

개발자는 jax-js를 활용해 소비자용 하드웨어에서 컴퓨팅 집약적인 애플리케이션을 로컬에서 빠르게 실행할 수 있다.