- 구글 DeepMind가 발표한, 모델을 읽기 쉽고 기능적인 Pytree 데이터 구조로 작성하기 위한 JAX 라이브러리
- 모델이 학습된 후에도 쉽게 작업할 수 있도록 하는 데 중점을 둠
- 모델 구성 요소를 리버스 엔지니어링 또는 제거 및 내부 활성화 검사 및 조사, 모델 수술, 아키텍처 디버깅 등과 관련된 연구에 적합
- 모델 시각화, 수정 및 분석용 도구를 제공해서, 그냥 모델을 만들고 훈련하는 용도로도 사용 가능
- 함께 설계되었지만 각각 독립적으로 사용 가능한 모듈형 도구 모음으로 구성
포함된 도구들
penzai.nn (pz.nn)
- Flax, Haiku, Keras, Equinox 등의 다른 신경망 라이브러리에 대한 대안으로, 선언적 컴비네이터 기반 신경망 라이브러리
- 모델의 forward pass 전체 구조를 모델 pytree에 노출하여, pretty printing으로 모델의 모든 작업을 볼 수 있고
jax.tree_util로 새로운 런타임 로직을 주입할 수 있음
- 모델은 내부적으로 호출 가능한 pytree에 불과함
penzai.treescope (pz.ts)
- 일반 IPython/Colab 렌더러를 대체하는 드롭인 방식으로 작동하는 강력한 대화형 Python pretty-printer
- Penzai 모델과 기타 깊이 중첩된 JAX pytree를 이해하는 데 도움이 되도록 설계되었으며, 임의 차원의 NDArray 시각화를 기본적으로 지원함
penzai.core.selectors (pz.select)
- JAX의
.at[...].set(...) 구문을 임의의 유형 기반 pytree 순회로 일반화하는 pytree 용 다기능 도구
- Penzai 모델 및 기타 데이터 구조의 복잡한 재작성 또는 온더플라이 패치를 쉽게 수행할 수 있음
penzai.core.named_axes (pz.nx)
- 일반 JAX 함수를 명명된 축에 대해 벡터화하고, 새로운 배열 API를 배우지 않고도 명명된 프로그래밍 스타일과 위치 프로그래밍 스타일 간에 원활하게 전환할 수 있는 경량 명명된 축 시스템
penzai.data_effects (pz.de)
- pytree 순회를 기반으로 구축되고 모델 작성이나 사용을 방해하지 않으면서 제어권을 갖는, 부수 인수, 난수 및 상태 변수에 대한 옵트인 시스템