What is JAX?

What is the common confusion about JAX?

JAX does not prohibit all stateful programming, but hidden mutation inside transformed functions can produce incorrect or unsupported behavior.

Learn JAX in the course

Start with

• Introduction to JAX

  PyTorch mutates tensors. TensorFlow builds graphs. JAX compiles pure functions. That last one changes how you think about deep learning.

  Phase 03: Deep Learning Core

Lessons that name JAX in a title or section

• Introduction to PyTorch

  You built the engine from pistons and crankshafts. Now learn the one everyone actually drives. Build and train neural networks using PyTorch's nn.Module, nn.Sequential, and autograd.

  Phase 03: Deep Learning Core

Taught in Phase 03: Deep Learning Core.

Related terms

• AutogradA system that records or transforms tensor operations so it can compute derivatives, usually with reverse-mode automatic differentiation.
• TensorA typed array with a shape, data type, and device placement that frameworks use to represent inputs, parameters, activations, and gradients.
• CUDANVIDIA's platform and programming model for general-purpose computation on compatible GPUs.

Sources

• JAX documentation

More terms in Math & training

• Activation Checkpointing
• Activation Function
• Adam (Optimizer)
• AdamW
• Autograd
• Backpropagation
• Batch Size
• Contrastive Learning
• Cross-Entropy
• Data Augmentation
• DPO (Direct Preference Optimization)
• Dropout
• Eigenvalue
• Epoch
• Fine-tuning
• Gradient
• Gradient Accumulation
• Gradient Clipping
• Gradient Descent
• Hyperparameter
• Knowledge Distillation
• Learning Rate
• Learning Rate Schedule
• LoRA (Low-Rank Adaptation)
• Loss Function
• Mixed Precision
• NaN (Not a Number)
• Normalization
• Optimizer
• Overfitting
• QLoRA
• ReLU
• RLHF (Reinforcement Learning from Human Feedback)
• SFT (Supervised Fine-Tuning)
• Softmax
• Stochastic Gradient Descent (SGD)
• Transfer Learning
• Underfitting
• Warmup
• Weight
• Weight Decay

Open the Math & training list in the glossary