LLM 시대 모두가 기다려온 확장 가능한 고성능 딥러닝 라이브러리 JAX/Flax
파이토치는 쓸 만한 라이브러리였다. LLM 전성시대가 닥치기 전까지는. JAX는 파이토치가 부족했던 부분을 채우며 부상했다. 대규모 계산의 병렬처리 등 바로 이 시대에 무엇보다 우선시되는 ‘확장성’을 염두에 두고 설계되었기 때문이다.
이 책은 국내 JAX 전문가들이 모인 모두의연구소 JAX/Flax LAB이 집필한 최초의 JAX+Flax 도서다. JAX 기초부터 시작해 함수형 프로그래밍, JIT 컴파일, 병렬처리 등 JAX의 특장점을 살펴본 다음, 현재 JAX와 가장 많이 조합되는 신경망 라이브러리인 Flax를 설명한다. CNN, ResNet, DCGAN, CLIP, DistilGPT2 모델의 관련 이론을 간단히 살펴보며 JAX와 Flax를 조합해서 우아하게 구현하는 방법을 보여준다. 새 술은 새 부대에. 파이토치는 놓아주고 새로운 시대에 맞는 새로운 라이브러리를 익혀보자.

주요 내용함수형 프로그래밍, 파이썬 라이브러리 등 JAX 사용 시 알아야 할 기초JIT 컴파일, 자동 벡터화, pytree, 병렬처리 등 JAX의 주요 특징CNN 튜토리얼로 알아보는 Flax 기초ResNet, DCGAN, CLIP 모델을 구축하며 Flax에 익숙해지기코랩, 캐글에서 TPU 환경 설정하기