Equinox

JAX 위에서 동작하는 라이브러리로, PyTorch와 유사한 객체 지향적 인터페이스를 제공하면서 JAX의 함수형 프로그래밍 원칙과 불변성을 유지하도록 돕는다. 모델 구조를 PyTree로 관리하여 JAX의 변환 기능들과 쉽게 호환된다.