Penzai - Toolkit JAX untuk membangun, mengedit, dan memvisualisasikan jaringan saraf

• Library JAX yang diumumkan oleh Google DeepMind untuk menulis model dalam struktur data Pytree yang fungsional dan mudah dibaca
• Berfokus untuk memudahkan pekerjaan dengan model bahkan setelah model selesai dilatih
  • Cocok untuk riset yang terkait dengan reverse engineering atau penghapusan komponen model, serta inspeksi dan investigasi aktivasi internal, bedah model, debugging arsitektur, dan lain-lain
  • Menyediakan alat untuk visualisasi, modifikasi, dan analisis model, sehingga juga bisa digunakan sekadar untuk membuat dan melatih model
• Terdiri dari kumpulan alat modular yang dirancang bersama, tetapi masing-masing dapat digunakan secara independen

Alat yang disertakan

penzai.nn (pz.nn)

• Library jaringan saraf berbasis combinator deklaratif sebagai alternatif untuk library jaringan saraf lain seperti Flax, Haiku, Keras, dan Equinox
• Mengekspos seluruh struktur forward pass model ke pytree model, sehingga semua operasi model dapat dilihat melalui pretty printing dan logika runtime baru dapat disuntikkan dengan jax.tree_util
• Model pada dasarnya hanyalah pytree yang dapat dipanggil

penzai.treescope (pz.ts)

• Pretty-printer Python interaktif yang kuat dan bekerja sebagai pengganti drop-in untuk renderer IPython/Colab biasa
• Dirancang untuk membantu memahami model Penzai dan pytree JAX lain yang bertingkat dalam, serta secara bawaan mendukung visualisasi NDArray berdimensi arbitrer

penzai.core.selectors (pz.select)

• Alat serbaguna untuk pytree yang menggeneralisasi sintaks .at[...].set(...) di JAX menjadi traversal pytree berbasis tipe yang arbitrer
• Memungkinkan rewrite kompleks atau patch on-the-fly pada model Penzai dan struktur data lainnya dengan mudah

penzai.core.named_axes (pz.nx)

• Sistem sumbu bernama yang ringan untuk memvektorkan fungsi JAX biasa terhadap sumbu bernama, serta memungkinkan perpindahan mulus antara gaya pemrograman bernama dan gaya pemrograman posisional tanpa perlu mempelajari API array baru

penzai.data_effects (pz.de)

• Sistem opt-in untuk argumen sampingan, bilangan acak, dan variabel status yang dibangun di atas traversal pytree serta tetap memberikan kontrol tanpa mengganggu penulisan atau penggunaan model

[brainer]

Apakah ada yang serupa untuk PyTorch?