blackjax-devs/blackjax

해결하는 문제

BlackJAX는 확률 프로그래밍 언어(PPL)에서 분리된 베이지안 추론을 위한 모듈식 샘플러 라이브러리를 제공합니다. 샘플링 알고리즘이 너무 단일화되거나 특정 모델링 프레임워크에 묶여 있는 문제를 해결하며, 사용자가 전체 PPL 없이도 CPU, GPU, TPU에서 샘플링을 수행할 수 있도록 합니다.

작동 방식

이 라이브러리는 상태에 의존하지 않는 함수가 현재 상태를 받아 새로운 상태와 전이 정보를 반환하는 일반적인 전이 커널 패턴을 기반으로 합니다. 일관된 API를 통해 사용자는 적분기, 제안 함수, 운동량 생성기와 같은 기본 블록을 쉽게 조합하고 교체하며, 복잡한 샘플링 알고리즘을 구축할 수 있습니다.

대상 사용자

• 새로운 샘플링 알고리즘을 개발하는 연구자 및 개발자.
• 안정적이고 성능이 뛰어난 샘플러 구현이 필요한 PPL 개발자.
• 로그 확률 밀도 함수(logpdf)가 있으며 단순히 샘플을 생성하기만 원하는 사용자.
• GPU/TPU를 활용해 하드웨어 가속을 원하는 사용자.

주요 특징

• 하드웨어 가속: JAX를 통해 CPU, GPU, TPU에서 작동 가능.
• 조합 가능한 설계: 전이 커널의 모듈식 아키텍처를 사용해 쉽게 사용자 정의하고 새로운 알고리즘을 생성할 수 있음.
• PPL 독립: JAX 호환 로그 확률 밀도 함수를 제공할 수 있는 모든 PPL과 통합 가능.
• 연구용 최적화: 새로운 추론 방식 개발 및 테스트를 가속화하기 위해 저수준 빌딩 블록을 공개.