JAX transformation, definícia a použitie | AI slovník

Odborná definícia

JAX transformation je operátor vyššieho rádu, ktorý prijme Python funkciu a vytvorí novú funkciu s upravenou výpočtovou sémantikou. Medzi základné transformácie patria grad pre diferenciáciu, jit pre kompiláciu, vmap pre automatické batchovanie a pmap alebo novšie shardované mechanizmy pre paralelizáciu. Transformácia funguje trasovaním primitívnych operácií do medzireprezentácie a ich interpretáciou podľa pravidiel transformácie. Skladateľnosť umožňuje napríklad diferencovať vektorizovanú funkciu, no vyžaduje čisté funkcie, statické tvary v relevantných miestach a kompatibilné riadenie toku.

Zrozumiteľné vysvetlenie

Namiesto prepisovania algoritmu pre gradient, batch a akcelerátor napíšete jednu čistú funkciu. Potom ju obalíte grad, vmap alebo jit a JAX vytvorí variant s novým správaním. Je to podobné, ako keby ste mali jeden recept a nástroje, ktoré z neho automaticky pripravia recept pre sto porcií, spočítajú citlivosť na každú ingredienciu alebo optimalizujú poradie krokov. Transformácie sa dajú skladať, ale iba vtedy, keď funkcia nepoužíva skryté zmeny stavu, ktoré trasovací mechanizmus nevidí.

Časté otázky

Čo znamená, že transformácie sú skladateľné?

Výstup jednej transformácie je opäť funkcia, ktorú možno odovzdať ďalšej. Napríklad jit(vmap(grad(f))) môže kompilovať batched gradienty. Poradie však ovplyvňuje tvary, pamäť a výkon, preto nie sú všetky kompozície prakticky ekvivalentné.

Prečo transformácia zlyhá na Python if?

Ak podmienka závisí od hodnoty známej až počas výpočtu na zariadení, Python chce rozhodnúť počas trasovania, keď má iba abstraktný tracer. Používa sa jax.lax.cond, maskovanie alebo statický argument podľa charakteru vetvenia.

Čo je tracer?

Tracer je abstraktný zástupca hodnoty, ktorý nesie informáciu potrebnú pre konkrétnu transformáciu, napríklad tvar a typ alebo tangentnú zložku. Operácie nad tracermi sa zaznamenávajú a nepredstavujú vždy okamžite vypočítané čísla.

Dá sa transformovať funkcia s vedľajším efektom?

Tlač, zápis do globálnej premennej alebo mutácia Python objektu sa pri trasovaní nemusí správať podľa očakávania a môže sa vykonať iba pri kompilácii. Stav sa preto odovzdáva explicitne a diagnostika používa nástroje určené pre JAX runtime.

Ako sa vytvára vlastné derivačné pravidlo?

JAX poskytuje custom_jvp a custom_vjp, ktorými sa určí dopredné alebo spätné pravidlo pre funkciu. Používa sa to pri numericky stabilnej derivácii, volaní externého primitíva alebo pri zámere zmeniť gradient bez zmeny dopredného výpočtu.

Súvisiace 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.