optax

JAX를 위한 경사 하강법 기반 최적화 라이브러리이다. Adam, RMSProp 등 다양한 최적화 알고리즘과 학습률 스케줄링 기능을 제공하며, 함수형 구조를 통해 신경망 학습 과정을 정밀하게 제어할 수 있다.