ml python jax

JAX – NumPy na sterydach

• Automatyczne różniczkowanie
• GPU i TPU
• Kompilacja JIT
• inny styl programowania (funkcyjny, immutable arrays).

JAX provides NumPy-style array operations through jax.numpy , which is commonly imported as jnp. You can create arrays, perform arithmetic, reshape data, calculate totals or averages, and carry out matrix operations using a similar style.

The important point is that JAX does not begin with neural networks.It begins with numerical calculations performed on arrays.

JAX	Feature Purpose
jax.grad()	Calculates gradients automatically
jax.jit()	Compiles a function for efficient execution
jax.vmap()	Applies a function across batches of data

JAX lets you write numerical functions and then transform those functions for different computational purposes.

JAX can be used for many types of numerical and scientific computing. It is not limited to machine learning.

Konersja z numpy:

# NumPy -> JAX
jarr = jnp.array(arr)