Keras

Keras 3 compiles a single model definition into executable code for JAX, TensorFlow, PyTorch, or OpenVINO by deferring all numerical operations to a pluggable backend abstraction layer. The active backend is selected at import time via KERAS_BACKEND environment variable or ~/.keras/keras.json and cannot be changed post-import. During model construction, symbolic execution via compute_output_spec() infers shapes and dtypes without computation; during training/inference, calls dispatch to backend-specific implementations in keras/src/backend/{jax,torch,tensorflow,openvino}/. This architecture enables write-once-run-anywhere model code without backend-specific rewrites.

Keras provides two high-level APIs for composing neural networks: Sequential (linear stack of layers) and Functional (arbitrary directed acyclic graphs with multiple inputs/outputs). Both APIs accept layer instances (Dense, Conv2D, LSTM, etc.) and automatically handle tensor shape inference, weight initialization, and forward pass construction. The Functional API supports layer sharing, multi-branch architectures, and residual connections by explicitly passing tensors between layer calls. Under the hood, layers inherit from keras.layers.Layer, which implements __call__ to dispatch to backend-specific compute_output_spec (symbolic) and call (eager) methods, enabling shape validation before execution.

Keras provides model.save() and keras.saving.load_model() for serializing and deserializing models. Models can be saved in three formats: Keras format (HDF5 or ZIP with architecture + weights), SavedModel (TensorFlow format with concrete functions), or ONNX. The Keras format stores model architecture as JSON and weights as HDF5 or NumPy files. Deserialization reconstructs the model from saved architecture and weights, and custom layers/losses/metrics can be registered via custom_objects parameter. Model checkpointing during training is handled by keras.callbacks.ModelCheckpoint, which saves the best model based on validation metrics. Weights can be saved/loaded independently via model.save_weights() and model.load_weights().

Keras provides keras.optimizers.schedules for learning rate scheduling (ExponentialDecay, CosineDecay, PolynomialDecay, etc.) and keras.callbacks for hyperparameter tuning (LearningRateScheduler, ReduceLROnPlateau). Learning rate schedules decay the learning rate over training steps or epochs to improve convergence. Callbacks enable dynamic hyperparameter adjustment during training (e.g., reducing learning rate when validation loss plateaus). Keras also integrates with external hyperparameter optimization frameworks (Keras Tuner, Optuna, Ray Tune) via callbacks. The fit() method accepts learning rate schedules and callbacks, enabling end-to-end hyperparameter optimization without custom training loops.

Keras enables custom layer implementation by subclassing keras.layers.Layer and implementing build() (weight initialization), call() (forward pass), and compute_output_spec() (shape inference). Custom loss functions can be implemented by subclassing keras.losses.Loss or as callables. Custom layers and losses automatically support automatic differentiation through the active backend (JAX, PyTorch, TensorFlow) without requiring manual gradient implementation. Custom operations can use keras.ops for backend-agnostic computation or backend-specific ops for optimization. The framework handles gradient computation, mixed-precision scaling, and distributed training for custom layers/losses without user code changes.

Keras provides model.summary() to print a human-readable summary of model architecture (layer names, output shapes, parameter counts, connectivity). The summary includes total trainable and non-trainable parameters, enabling quick model size estimation. Keras also supports model graph visualization via keras.utils.plot_model(), which generates a visual diagram of the model architecture (useful for Functional API models with complex connectivity). Model introspection methods (model.get_config(), model.get_weights()) enable programmatic access to architecture and weights. These tools are backend-agnostic and work identically across JAX, PyTorch, and TensorFlow.

Keras provides built-in regularization through layer parameters and dedicated layers: kernel_regularizer/bias_regularizer (L1/L2 weight regularization), activity_regularizer (activation regularization), Dropout layer (random unit dropping), and BatchNormalization layer (feature normalization with learnable scale/shift). Regularization is applied during training via the loss function (for weight regularization) or forward pass (for dropout, batch norm). Dropout randomly zeros activations during training and scales them during inference. BatchNormalization normalizes activations to zero mean and unit variance, reducing internal covariate shift and enabling higher learning rates. All regularization techniques are backend-agnostic and work identically across JAX, PyTorch, and TensorFlow.

Keras delegates automatic differentiation to the active backend (JAX's jax.grad, PyTorch's autograd, TensorFlow's tf.GradientTape) through a unified keras.ops interface that wraps backend-specific gradient functions. During training, the fit() method constructs a loss function, computes gradients via backend-native autodiff, and applies optimizer updates. Custom training loops can use keras.ops.grad() to compute gradients of arbitrary functions. The backend abstraction ensures that gradient computation, mixed-precision scaling, and gradient clipping work identically across JAX, PyTorch, and TensorFlow without user code changes.

Keras provides a comprehensive library of pre-implemented layers (Dense, Conv1D/2D/3D, LSTM, GRU, Attention, BatchNormalization, Dropout, etc.) in keras.layers, each with configurable parameters (units, activation, regularization, initialization). Layers are backend-agnostic; their implementations in keras/src/layers/ use only keras.ops (NumPy-compatible operations) and backend-specific ops, ensuring portability across JAX, PyTorch, and TensorFlow. Each layer implements build() (weight initialization), call() (forward pass), and compute_output_spec() (shape inference). Custom layers can be created by subclassing keras.layers.Layer and implementing these methods.

Keras's fit() method provides a high-level training interface that handles gradient computation, optimizer updates, metric tracking, and validation in a single call. The method accepts a model, training data (NumPy arrays, tf.data.Dataset, or backend-native iterables), loss function, optimizer, and metrics. During training, fit() iterates over batches, computes loss and gradients via backend autodiff, applies optimizer updates, and accumulates metrics. Callbacks (keras.callbacks.Callback) hook into training events (epoch start/end, batch end) for logging, early stopping, learning rate scheduling, and checkpointing. Validation is performed at configurable intervals, and metrics are computed on both training and validation sets.

Keras exposes a NumPy-compatible operations API (keras.ops) that wraps backend-specific implementations (JAX, PyTorch, TensorFlow) for mathematical operations (matmul, reshape, concatenate, etc.), neural network operations (conv2d, batch_norm, etc.), and activation functions. Each operation in keras.ops has implementations in keras/src/ops/{numpy,nn,core}.py that dispatch to the active backend. This enables users to write backend-agnostic code using familiar NumPy-like syntax. Operations support automatic differentiation through backend autodiff, and the API includes both eager execution (immediate computation) and symbolic execution (shape/dtype inference via compute_output_spec).

Keras provides model.export() and backend-specific export functions to convert trained models into deployment-ready formats: SavedModel (TensorFlow), ONNX (cross-framework), LiteRT (mobile), and OpenVINO (edge inference). Export functions in keras/src/saving/ serialize model architecture, weights, and preprocessing layers into format-specific representations. SavedModel export includes a concrete function signature for inference. ONNX export converts Keras ops to ONNX operators via a mapping layer. LiteRT and OpenVINO exports optimize models for mobile and edge devices. Exported models can be loaded and used for inference without Keras, enabling deployment on diverse hardware (mobile, edge, cloud).

Keras supports distributed training via keras.distribution.DataParallel (data parallelism) and backend-specific distributed APIs (tf.distribute.Strategy for TensorFlow, torch.nn.DataParallel for PyTorch, jax.pmap for JAX). Data parallelism splits training data across devices, computes gradients on each device, and synchronizes gradients across devices before optimizer updates. The fit() method automatically handles distributed training when a distribution strategy is configured. Gradient synchronization and optimizer updates are coordinated by the distribution backend, ensuring convergence across devices. Keras abstracts distribution details, allowing the same model code to scale from single-GPU to multi-GPU/TPU without modification.

Keras supports quantization (reducing precision from float32 to int8/float16) and mixed-precision training (using float16 for computation, float32 for weights) to reduce memory usage and accelerate training. Quantization is implemented via keras.quantizers (post-training quantization) and quantization-aware training (QAT) layers. Mixed-precision training is enabled via keras.mixed_precision.set_global_policy(), which automatically casts operations to lower precision while maintaining numerical stability. The optimizer applies loss scaling to prevent gradient underflow in float16. Quantized models can be exported to optimized formats (LiteRT, OpenVINO) for deployment on resource-constrained devices.

Keras provides preprocessing layers (keras.layers.preprocessing.*) for common data transformations: image augmentation (RandomFlip, RandomRotation, RandomZoom), text preprocessing (TextVectorization, Hashing), and numerical feature engineering (Normalization, Discretization). Preprocessing layers are stateful (they learn statistics from training data via adapt()) and can be included in models for end-to-end training. During training, preprocessing is applied on-device (GPU/TPU) for efficiency. Preprocessing layers support both eager and symbolic execution, enabling shape inference and batch processing. Exported models include preprocessing layers, enabling end-to-end inference without external preprocessing code.