JAX

JAX provides array operations with a familiar numerical interface and transformations that act on functions. Automatic differentiation computes derivatives, just-in-time compilation prepares compatible calculations for execution and vectorization maps operations over batches. These transformations compose, allowing a numerical function to become a differentiated or compiled training computation. The programming model emphasizes explicit data flow and functional treatment of state, including random-number keys. Python control flow and side effects can behave differently when tracing or compiling a function. Competence therefore includes understanding which values are static, how shapes influence compiled programs and when asynchronous device execution affects measurement.

Write small pure numerical functions with explicit inputs and outputs, then add transformations while checking results against untransformed cases. Manage parameters, optimizer state and random keys explicitly. Inspect array shapes and device placement, and separate compilation time from steady-state execution when benchmarking. Use appropriate debugging tools for traced computations and verify numerical stability. The result should be a reproducible computation or training loop whose state transitions and transformations are clear, rather than code that happens to work eagerly but changes behavior when compiled or vectorized.

Tracing and compilation introduce constraints on Python behavior, shapes and mutation. Recompilation can dominate workloads with changing static inputs, and asynchronous execution can make naive timing misleading. Random-key reuse can create unintended dependence. Automatic gradients still differentiate the implemented objective, including mistakes in it. JAX differs from a complete high-level training framework; additional libraries may supply model and optimizer abstractions. Validate transformed computations and measure end-to-end workloads rather than assuming compilation always improves runtime.