Funkční transformace v JAX | AI slovník

Odborná definice

JAX transformation je operátor vyššieho rádu, který prijme Python funkci a vytvoří novou funkci s upravenou výpočtovou sémantikou. Mezi základní transformace patří grad pro diferenciaci, jit pro kompilaci, vmap pro automatické batchování a pmap nebo novější shardované mechanizmy pro paralelizaci. Transformace funguje trasováním primitívnych operací do medzireprezentace a jejich interpretací podle pravidel transformace. Skladatelnost umožňuje například diferencovat vektorizovanou funkci, ale vyžaduje čisté funkce, statické tvary v relevantních miestach a kompatibilní řízení toku.

Srozumitelné vysvětlení

Místo prepisování algoritmu pro gradient, batch a akcelerátor napíšete jednu čistou funkci. Potom ji obalíte grad, vmap nebo jit a JAX vytvoří variant s novým chováním. Je to podobné, jako keby ste měli jeden recept a nástroje, které z něho automaticky pripravia recept pro sto porcií, sečtou citlivost na každou ingredienci nebo optimalizují pořadí kroků. Transformace se dají skladat, ale pouze tehdy, když funkce nepoužívá skryté změny stavu, které trasovací mechanismus nevidí.

Časté otázky

Co znamená, že transformace jsou skladatelné?

Výstup jedné transformace je opeť funkce, kterou lze odovzdat další. Například jit(vmap(grad(f))) může kompilovat batched gradienty. Pořadí však ovlivňuje tvary, paměť a výkon, proto nejsou všechny kompozíce prakticky ekvivalentné.

Proč transformace selže na Python if?

Pokud podmínka závisí na hodnoty známej až během výpočtu na zařízení, Python chce rozhodnout během trasování, když má pouze abstraktný tracer. Používá se jax.lax.cond, maskování nebo statický argument podle charakteru vetvení.

Co je tracer?

Tracer je abstraktný zástupca hodnoty, který nese informaci potřebnou pro konkrétní transformaci, například tvar a typ nebo tangentnou zložku. Operace nad tracermi se zaznamenávají a nepredstavují vždy okamžitě vypočítané čísla.

Lze transformovat funkce s vedlajším efektom?

Tlač, zápis do globální proměnné nebo mutace Python objektu se při trasování nemusí správat podle očakávání a může se provést pouze při kompilaci. Stav se proto odovzdáva explicitně a diagnostika používá nástroje určené pro JAX runtime.

Jak se vytváří vlastní derivačné pravidlo?

JAX poskytuje custom_jvp a custom_vjp, kterými se určí dopredné nebo zpětné pravidlo pro funkci. Používá se to při numericky stabilní derivaci, volaní externího primitíva nebo při zámere změnit gradient bez změny dopredného výpočtu.

Související pojmy

Zdroje a redakční stopa

JAX, Automatic differentiation JAX, Automatic vectorization JAX, Autodidax

Definícia je autorská odborná syntéza. Pri právnych a regulačných rozhodnutiach má prednosť aktuálne oficiálne znenie predpisu a posúdenie konkrétneho prípadu.