Glossary term
Glossary term
Infrastructure and Serving
A function that takes a function as input and returns a transformed function as output. JAX uses function transformations.
G
Created for this library
An ML researcher uses JAX function transformations like jit and grad to compile and differentiate a model written in pure Python.
A research engineer uses pmap as a function transformation to scale a JAX training step across multiple devices with minimal code change.
A research team uses vmap as a function transformation to batch a sample-level function so it runs efficiently over a whole batch.
Definition source: Google for Developers Machine Learning Glossary | Creative Commons Attribution 4.0 License